Add support for OIDC ID token authentication using an environment variables and files

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;

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<NamedIDTokenSource> 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",

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);
+    }
+  }
+}

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:
+   *     <ul>
+   *       <li>If the file path is invalid or malformed
+   *       <li>If the file does not exist
+   *       <li>If there are security permission issues accessing the file
+   *       <li>If the file is empty or contains only whitespace
+   *       <li>If the file cannot be read due to I/O errors
+   *       <li>If the token format in the file is invalid
+   *     </ul>
+   */
+  @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<String> 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<String> 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));
+    }
+  }
+}

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()) {

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<String, String> envVars) {
+    return new Environment(envVars, new String[0], "test");
+  }
+
+  private static Stream<Arguments> 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<String, String> createEnvVars(String key, String value) {
+    Map<String, String> envVars = new HashMap<>();
+    envVars.put(key, value);
+    return envVars;
+  }
+
+  @ParameterizedTest(name = "{0}")
+  @MethodSource("provideTestCases")
+  void testGetIDToken(
+      String testName,
+      String envVarName,
+      Map<String, String> envVars,
+      String expectedToken,
+      Class<? extends Exception> 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());
+    }
+  }
+}