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

databricks-sdk-java/src/main/java/com/databricks/sdk/core/oauth/EnvVarIDTokenSource.java

@@ -0,0 +1,52 @@
+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, or the
+   *     environment is null.
+   * @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");
+    }
+
+    if (env == null) {
+      throw new IllegalArgumentException("Environment cannot be null");
+    }
+
+    try {
+      String token = env.get(envVarName);
+      return new IDToken(token);
+    } catch (IllegalArgumentException e) {
+      throw new DatabricksException(
+          "Received empty ID token from environment variable " + envVarName);
+    }
+  }
+}

databricks-sdk-java/src/main/java/com/databricks/sdk/core/oauth/FileIDTokenSource.java

@@ -0,0 +1,71 @@
+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.Path;
+import java.nio.file.Paths;
+import java.util.List;
+
+/**
+ * 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.
+ *
+ * @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.
+   *
+   * @param audience The intended recipient of the ID Token (not used).
+   * @return An {@link IDToken} containing the token value from the file.
+   * @throws IllegalArgumentException if the file path is null or empty.
+   * @throws DatabricksException if the file does not exist, is empty, or contains only whitespace.
+   */
+  @Override
+  public IDToken getIDToken(String audience) {
+    if (Strings.isNullOrEmpty(filePath)) {
+      throw new IllegalArgumentException("File path cannot be null or empty");
+    }
+
+    try {
+      Path path = Paths.get(filePath);
+      if (!Files.exists(path)) {
+        throw new DatabricksException("File " + filePath + " does not exist");
+      }
+
+      List<String> lines = Files.readAllLines(path, StandardCharsets.UTF_8);
+      if (lines.isEmpty()) {
+        throw new DatabricksException("File " + filePath + " is empty");
+      }
+
+      try {
+        String token = lines.get(0).trim();
+        return new IDToken(token);
+      } catch (IllegalArgumentException e) {
+        throw new DatabricksException("Received empty ID token from file " + filePath);
+      }
+
+    } catch (IOException e) {
+      throw new DatabricksException(
+          "Failed to read ID token from file " + filePath + ": " + e.getMessage(), e);
+    }
+  }
+}

databricks-sdk-java/src/test/java/com/databricks/sdk/core/oauth/EnvVarIDTokenSourceTest.java

@@ -0,0 +1,91 @@
+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),
+        // Test case: Null environment
+        Arguments.of(
+            "Null environment", TEST_ENV_VAR_NAME, null, null, IllegalArgumentException.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());
+    }
+  }
+}

databricks-sdk-java/src/test/java/com/databricks/sdk/core/oauth/FileIDTokenSourceTest.java

@@ -0,0 +1,72 @@
+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 Path createTestFile(String content) throws IOException {
+    Path file = tempDir.resolve("test-token.txt");
+    Files.write(file, content.getBytes(StandardCharsets.UTF_8));
+    return file;
+  }
+
+  private static Stream<Arguments> provideTestCases() {
+    return Stream.of(
+        // Test case name, file path (null means create temp file), file content, expected token,
+        // expected exception
+        Arguments.of("Valid token file", null, TEST_TOKEN, TEST_TOKEN, null),
+        Arguments.of("Token with whitespace", null, "  " + TEST_TOKEN + "  ", TEST_TOKEN, null),
+        Arguments.of("Empty file", null, "", null, DatabricksException.class),
+        Arguments.of("File with only whitespace", null, "   ", null, DatabricksException.class),
+        Arguments.of("Null file path", null, null, null, IllegalArgumentException.class),
+        Arguments.of("Empty file path", "", null, null, IllegalArgumentException.class),
+        Arguments.of(
+            "Non-existent file",
+            "/path/to/nonexistent/file.txt",
+            null,
+            null,
+            DatabricksException.class));
+  }
+
+  @ParameterizedTest(name = "{0}")
+  @MethodSource("provideTestCases")
+  void testGetIDToken(
+      String testName,
+      String filePath,
+      String fileContent,
+      String expectedToken,
+      Class<? extends Exception> expectedException)
+      throws IOException {
+    // If filePath is null, create a temporary test file
+    String actualFilePath = filePath;
+    if (filePath == null && fileContent != null) {
+      actualFilePath = createTestFile(fileContent).toString();
+    }
+
+    FileIDTokenSource source = new FileIDTokenSource(actualFilePath);
+
+    if (expectedException != null) {
+      assertThrows(expectedException, () -> source.getIDToken(TEST_AUDIENCE));
+    } else {
+      IDToken token = source.getIDToken(TEST_AUDIENCE);
+      assertNotNull(token);
+      assertEquals(expectedToken, token.getValue());
+    }
+  }
+}