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 fcb79c87b..98d75d4bc 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 @@ -148,6 +148,14 @@ public class DatabricksConfig { @ConfigAttribute(env = "TOKEN_AUDIENCE") private String tokenAudience; + /** Path to the file containing an OIDC ID token. */ + @ConfigAttribute(env = "DATABRICKS_OIDC_TOKEN_FILEPATH", auth = "file-oidc") + private String oidcTokenFilepath; + + /** Environment variable name that contains an OIDC ID token. */ + @ConfigAttribute(env = "DATABRICKS_OIDC_TOKEN_ENV", auth = "env-oidc") + private String oidcTokenEnv; + public Environment getEnv() { return env; } @@ -528,6 +536,24 @@ public DatabricksConfig setTokenAudience(String tokenAudience) { return this; } + public String getOidcTokenFilepath() { + return oidcTokenFilepath; + } + + public DatabricksConfig setOidcTokenFilepath(String oidcTokenFilepath) { + this.oidcTokenFilepath = oidcTokenFilepath; + return this; + } + + public String getOidcTokenEnv() { + return oidcTokenEnv; + } + + public DatabricksConfig setOidcTokenEnv(String oidcTokenEnv) { + this.oidcTokenEnv = oidcTokenEnv; + return this; + } + public boolean isAzure() { if (azureWorkspaceResourceId != null) { return true; diff --git a/databricks-sdk-java/src/main/java/com/databricks/sdk/core/DefaultCredentialsProvider.java b/databricks-sdk-java/src/main/java/com/databricks/sdk/core/DefaultCredentialsProvider.java index c3fc3b1e4..0e4723f36 100644 --- a/databricks-sdk-java/src/main/java/com/databricks/sdk/core/DefaultCredentialsProvider.java +++ b/databricks-sdk-java/src/main/java/com/databricks/sdk/core/DefaultCredentialsProvider.java @@ -1,6 +1,7 @@ package com.databricks.sdk.core; import com.databricks.sdk.core.oauth.*; +import com.google.common.base.Strings; import java.util.ArrayList; import java.util.List; import org.slf4j.Logger; @@ -111,6 +112,20 @@ private void addOIDCCredentialsProviders(DatabricksConfig config) { } List namedIdTokenSources = new ArrayList<>(); + namedIdTokenSources.add( + new NamedIDTokenSource( + "env-oidc", + new EnvVarIDTokenSource( + // Use configured environment variable name if set, otherwise default to + // DATABRICKS_OIDC_TOKEN + Strings.isNullOrEmpty(config.getOidcTokenEnv()) + ? "DATABRICKS_OIDC_TOKEN" + : config.getOidcTokenEnv(), + config.getEnv()))); + + namedIdTokenSources.add( + new NamedIDTokenSource("file-oidc", new FileIDTokenSource(config.getOidcTokenFilepath()))); + namedIdTokenSources.add( new NamedIDTokenSource( "github-oidc", diff --git a/databricks-sdk-java/src/main/java/com/databricks/sdk/core/oauth/EnvVarIDTokenSource.java b/databricks-sdk-java/src/main/java/com/databricks/sdk/core/oauth/EnvVarIDTokenSource.java new file mode 100644 index 000000000..5bbc1c4e2 --- /dev/null +++ b/databricks-sdk-java/src/main/java/com/databricks/sdk/core/oauth/EnvVarIDTokenSource.java @@ -0,0 +1,47 @@ +package com.databricks.sdk.core.oauth; + +import com.databricks.sdk.core.DatabricksException; +import com.databricks.sdk.core.utils.Environment; +import com.google.common.base.Strings; + +/** Implementation of {@link IDTokenSource} that reads the ID token from an environment variable. */ +public class EnvVarIDTokenSource implements IDTokenSource { + /* The name of the environment variable to read the ID token from. */ + private final String envVarName; + /* The environment to read variables from. */ + private final Environment env; + + /** + * Creates a new EnvVarIDTokenSource that reads from the specified environment variable. + * + * @param envVarName The name of the environment variable to read the ID token from. + * @param env The environment to read variables from. + */ + public EnvVarIDTokenSource(String envVarName, Environment env) { + this.envVarName = envVarName; + this.env = env; + } + + /** + * Retrieves an ID Token from the environment variable. + * + * @param audience The intended recipient of the ID Token (unused in this implementation). + * @return An {@link IDToken} containing the token value from the environment variable. + * @throws IllegalArgumentException if the environment variable name is null or empty. + * @throws DatabricksException if the environment variable is not set or is empty. + */ + @Override + public IDToken getIDToken(String audience) { + if (Strings.isNullOrEmpty(envVarName)) { + throw new IllegalArgumentException("Environment variable name cannot be null or empty"); + } + + try { + String token = env.get(envVarName); + return new IDToken(token); + } catch (IllegalArgumentException e) { + throw new DatabricksException( + String.format("Received empty ID token from environment variable %s", envVarName), e); + } + } +} diff --git a/databricks-sdk-java/src/main/java/com/databricks/sdk/core/oauth/FileIDTokenSource.java b/databricks-sdk-java/src/main/java/com/databricks/sdk/core/oauth/FileIDTokenSource.java new file mode 100644 index 000000000..b503d454c --- /dev/null +++ b/databricks-sdk-java/src/main/java/com/databricks/sdk/core/oauth/FileIDTokenSource.java @@ -0,0 +1,123 @@ +package com.databricks.sdk.core.oauth; + +import com.databricks.sdk.core.DatabricksException; +import com.google.common.base.Strings; +import java.io.IOException; +import java.nio.charset.StandardCharsets; +import java.nio.file.Files; +import java.nio.file.InvalidPathException; +import java.nio.file.Path; +import java.nio.file.Paths; +import java.util.List; +import java.util.stream.Collectors; + +/** + * Implementation of {@link IDTokenSource} that reads the ID token from a file. The token is read + * using UTF-8 encoding and any leading/trailing whitespace is trimmed. The file should contain + * exactly one non-empty line with the token value. Files with multiple non-empty lines or only + * empty lines will result in an error. + * + * @see IDTokenSource + */ +public class FileIDTokenSource implements IDTokenSource { + /* The path to the file containing the ID token. */ + private final String filePath; + + /** + * Creates a new FileIDTokenSource that reads from the specified file. + * + * @param filePath Path to the file containing the ID token. The file should contain a single line + * with the token value. + * @throws IllegalArgumentException if the file path is null or empty. + */ + public FileIDTokenSource(String filePath) { + this.filePath = filePath; + } + + /** + * Retrieves an ID Token from the file. The file is read using UTF-8 encoding and the first line + * is used as the token value. Any leading or trailing whitespace in the token is trimmed. + * + * @param audience The intended recipient of the ID Token. This parameter is not used in this + * implementation as the token is read directly from the file. + * @return An {@link IDToken} containing the token value from the file. + * @throws IllegalArgumentException if the file path is null or empty. + * @throws DatabricksException in the following cases: + * + */ + @Override + public IDToken getIDToken(String audience) { + if (Strings.isNullOrEmpty(filePath)) { + throw new IllegalArgumentException("File path cannot be null or empty"); + } + + Path path; + try { + path = Paths.get(filePath); + } catch (InvalidPathException e) { + throw new DatabricksException("Invalid file path: " + filePath, e); + } + + boolean isFileExists; + try { + isFileExists = Files.exists(path); + } catch (SecurityException e) { + throw new DatabricksException( + String.format( + "Security permission denied when checking if file %s exists: %s", + filePath, e.getMessage()), + e); + } + + if (!isFileExists) { + throw new DatabricksException(String.format("File %s does not exist", filePath)); + } + + List rawLines; + try { + rawLines = Files.readAllLines(path, StandardCharsets.UTF_8); + } catch (IOException e) { + throw new DatabricksException( + String.format("Failed to read ID token from file %s: %s", filePath, e.getMessage()), e); + } catch (SecurityException e) { + throw new DatabricksException( + String.format( + "Security permission denied when reading file %s: %s", filePath, e.getMessage()), + e); + } + + // Filter out empty lines + List nonEmptyLines = + rawLines.stream() + .map(String::trim) + .filter(line -> !line.isEmpty()) + .collect(Collectors.toList()); + + if (nonEmptyLines.isEmpty()) { + throw new DatabricksException(String.format("File %s contains only empty lines", filePath)); + } + + if (nonEmptyLines.size() > 1) { + throw new DatabricksException( + String.format( + "The token should be a single line but the file %s contains %d non-empty lines", + filePath, nonEmptyLines.size())); + } + + String token = nonEmptyLines.get(0); + + try { + return new IDToken(token); + } catch (IllegalArgumentException e) { + throw new DatabricksException( + String.format("Received empty ID token from file %s", filePath)); + } + } +} diff --git a/databricks-sdk-java/src/main/java/com/databricks/sdk/core/oauth/IDToken.java b/databricks-sdk-java/src/main/java/com/databricks/sdk/core/oauth/IDToken.java index e9c3f1ac6..952fb6ed9 100644 --- a/databricks-sdk-java/src/main/java/com/databricks/sdk/core/oauth/IDToken.java +++ b/databricks-sdk-java/src/main/java/com/databricks/sdk/core/oauth/IDToken.java @@ -12,6 +12,7 @@ public class IDToken { * Constructs an IDToken with a value. * * @param value The ID Token string. + * @throws IllegalArgumentException if the token value is null or empty. */ public IDToken(String value) { if (value == null || value.isEmpty()) { diff --git a/databricks-sdk-java/src/test/java/com/databricks/sdk/core/oauth/EnvVarIDTokenSourceTest.java b/databricks-sdk-java/src/test/java/com/databricks/sdk/core/oauth/EnvVarIDTokenSourceTest.java new file mode 100644 index 000000000..49bf6536c --- /dev/null +++ b/databricks-sdk-java/src/test/java/com/databricks/sdk/core/oauth/EnvVarIDTokenSourceTest.java @@ -0,0 +1,88 @@ +package com.databricks.sdk.core.oauth; + +import static org.junit.jupiter.api.Assertions.*; + +import com.databricks.sdk.core.DatabricksException; +import com.databricks.sdk.core.utils.Environment; +import java.util.HashMap; +import java.util.Map; +import java.util.stream.Stream; +import org.junit.jupiter.params.ParameterizedTest; +import org.junit.jupiter.params.provider.Arguments; +import org.junit.jupiter.params.provider.MethodSource; + +/** Tests for EnvVarIDTokenSource. */ +public class EnvVarIDTokenSourceTest { + private static final String TEST_ENV_VAR_NAME = "TEST_ID_TOKEN"; + private static final String TEST_TOKEN = "test-id-token"; + private static final String TEST_AUDIENCE = "test-audience"; + + private Environment createTestEnvironment(Map envVars) { + return new Environment(envVars, new String[0], "test"); + } + + private static Stream provideTestCases() { + return Stream.of( + // Test case: Success case + Arguments.of( + "Success case", + TEST_ENV_VAR_NAME, + createEnvVars(TEST_ENV_VAR_NAME, TEST_TOKEN), + TEST_TOKEN, + null), + // Test case: Null environment variable name + Arguments.of( + "Null environment variable name", + null, + new HashMap<>(), + null, + IllegalArgumentException.class), + // Test case: Empty environment variable name + Arguments.of( + "Empty environment variable name", + "", + new HashMap<>(), + null, + IllegalArgumentException.class), + // Test case: Missing environment variable + Arguments.of( + "Missing environment variable", + TEST_ENV_VAR_NAME, + new HashMap<>(), + null, + DatabricksException.class), + // Test case: Empty token value + Arguments.of( + "Empty token value", + TEST_ENV_VAR_NAME, + createEnvVars(TEST_ENV_VAR_NAME, ""), + null, + DatabricksException.class)); + } + + private static Map createEnvVars(String key, String value) { + Map envVars = new HashMap<>(); + envVars.put(key, value); + return envVars; + } + + @ParameterizedTest(name = "{0}") + @MethodSource("provideTestCases") + void testGetIDToken( + String testName, + String envVarName, + Map envVars, + String expectedToken, + Class expectedException) { + Environment env = envVars != null ? createTestEnvironment(envVars) : null; + EnvVarIDTokenSource source = new EnvVarIDTokenSource(envVarName, env); + + if (expectedException != null) { + assertThrows(expectedException, () -> source.getIDToken(TEST_AUDIENCE)); + } else { + IDToken token = source.getIDToken(TEST_AUDIENCE); + assertNotNull(token); + assertEquals(expectedToken, token.getValue()); + } + } +} diff --git a/databricks-sdk-java/src/test/java/com/databricks/sdk/core/oauth/FileIDTokenSourceTest.java b/databricks-sdk-java/src/test/java/com/databricks/sdk/core/oauth/FileIDTokenSourceTest.java new file mode 100644 index 000000000..4565aafac --- /dev/null +++ b/databricks-sdk-java/src/test/java/com/databricks/sdk/core/oauth/FileIDTokenSourceTest.java @@ -0,0 +1,75 @@ +package com.databricks.sdk.core.oauth; + +import static org.junit.jupiter.api.Assertions.*; + +import com.databricks.sdk.core.DatabricksException; +import java.io.IOException; +import java.nio.charset.StandardCharsets; +import java.nio.file.Files; +import java.nio.file.Path; +import java.util.stream.Stream; +import org.junit.jupiter.api.io.TempDir; +import org.junit.jupiter.params.ParameterizedTest; +import org.junit.jupiter.params.provider.Arguments; +import org.junit.jupiter.params.provider.MethodSource; + +/** Tests for FileIDTokenSource. */ +public class FileIDTokenSourceTest { + private static final String TEST_TOKEN = "test-id-token"; + private static final String TEST_AUDIENCE = "test-audience"; + + @TempDir Path tempDir; + + private static Stream provideTestCases() { + return Stream.of( + // Test case name, fileContent, fileToReadFrom, expected token, expected exception + Arguments.of("Valid token file", TEST_TOKEN, "token.txt", TEST_TOKEN, null), + Arguments.of( + "Token with whitespace", " " + TEST_TOKEN + " ", "token.txt", TEST_TOKEN, null), + Arguments.of("Empty file", "", "token.txt", null, DatabricksException.class), + Arguments.of( + "File with only whitespace", " ", "token.txt", null, DatabricksException.class), + Arguments.of("Null file path", TEST_TOKEN, null, null, IllegalArgumentException.class), + Arguments.of("Empty file path", TEST_TOKEN, "", null, IllegalArgumentException.class), + Arguments.of( + "Non-existent file", TEST_TOKEN, "nonexistent.txt", null, DatabricksException.class)); + } + + @ParameterizedTest(name = "{0}") + @MethodSource("provideTestCases") + void testGetIDToken( + String testName, + String fileContent, + String fileToReadFrom, + String expectedToken, + Class expectedException) + throws IOException { + // Always create token.txt with the specified content + Path tokenFile = tempDir.resolve("token.txt"); + Files.write(tokenFile, fileContent.getBytes(StandardCharsets.UTF_8)); + + String testPathToRead = null; + // If fileToReadFrom is null, we want to simulate passing a null path to FileIDTokenSource (for + // error cases). + // If fileToReadFrom is an empty string, we want to simulate passing an empty path (also for + // error cases). + // Otherwise, resolve the file name relative to the temp directory to get the full path. + if (fileToReadFrom != null) { + if (fileToReadFrom.equals("")) { + testPathToRead = ""; + } else { + testPathToRead = tempDir.resolve(fileToReadFrom).toString(); + } + } + + FileIDTokenSource source = new FileIDTokenSource(testPathToRead); + + if (expectedException != null) { + assertThrows(expectedException, () -> source.getIDToken(TEST_AUDIENCE)); + } else { + IDToken token = source.getIDToken(TEST_AUDIENCE); + assertNotNull(token); + assertEquals(expectedToken, token.getValue()); + } + } +}