From ce06c215f0b51064f60bbb6859ba28e65ae285de Mon Sep 17 00:00:00 2001 From: Tanmay Rustagi Date: Thu, 22 Jan 2026 13:27:19 +0100 Subject: [PATCH 1/6] [Feature] CLI Profile support for unified host --- .../sdk/core/DatabricksCliCredentialsProvider.java | 13 ++++++++++++- 1 file changed, 12 insertions(+), 1 deletion(-) diff --git a/databricks-sdk-java/src/main/java/com/databricks/sdk/core/DatabricksCliCredentialsProvider.java b/databricks-sdk-java/src/main/java/com/databricks/sdk/core/DatabricksCliCredentialsProvider.java index e4376f203..57fb22e7c 100644 --- a/databricks-sdk-java/src/main/java/com/databricks/sdk/core/DatabricksCliCredentialsProvider.java +++ b/databricks-sdk-java/src/main/java/com/databricks/sdk/core/DatabricksCliCredentialsProvider.java @@ -31,7 +31,18 @@ private CliTokenSource getDatabricksCliTokenSource(DatabricksConfig config) { } List cmd = new ArrayList<>(Arrays.asList(cliPath, "auth", "token", "--host", config.getHost())); - if (config.getClientType() == ClientType.ACCOUNT) { + if (config.getExperimentalIsUnifiedHost() != null && config.getExperimentalIsUnifiedHost()) { + // For unified hosts, pass account_id, workspace_id, and experimental flag + cmd.add("--experimental-is-unified-host"); + if (config.getAccountId() != null) { + cmd.add("--account-id"); + cmd.add(config.getAccountId()); + } + if (config.getWorkspaceId() != null) { + cmd.add("--workspace-id"); + cmd.add(config.getWorkspaceId()); + } + } else if (config.getClientType() == ClientType.ACCOUNT) { cmd.add("--account-id"); cmd.add(config.getAccountId()); } From 42db1b9d2b07b569889449c44f2b5019b9f9988f Mon Sep 17 00:00:00 2001 From: Tanmay Rustagi Date: Thu, 22 Jan 2026 15:02:00 +0100 Subject: [PATCH 2/6] - --- .../DatabricksCliCredentialsProviderTest.java | 193 ++++++++++++++++++ 1 file changed, 193 insertions(+) create mode 100644 databricks-sdk-java/src/test/java/com/databricks/sdk/core/DatabricksCliCredentialsProviderTest.java diff --git a/databricks-sdk-java/src/test/java/com/databricks/sdk/core/DatabricksCliCredentialsProviderTest.java b/databricks-sdk-java/src/test/java/com/databricks/sdk/core/DatabricksCliCredentialsProviderTest.java new file mode 100644 index 000000000..67de6faf1 --- /dev/null +++ b/databricks-sdk-java/src/test/java/com/databricks/sdk/core/DatabricksCliCredentialsProviderTest.java @@ -0,0 +1,193 @@ +package com.databricks.sdk.core; + +import static org.junit.jupiter.api.Assertions.*; + +import com.databricks.sdk.core.oauth.Token; +import com.databricks.sdk.core.utils.Environment; +import java.lang.reflect.Method; +import java.time.Instant; +import java.util.List; +import org.junit.jupiter.api.Test; + +class DatabricksCliCredentialsProviderTest { + + private static final String HOST = "https://db-deco-test.databricks.com"; + private static final String ACCOUNT_ID = "968367da-7edd-44f7-9dea-3e0b20b0ec97"; + private static final String WORKSPACE_ID = "470576644108500"; + + /** + * Helper method to invoke the private getDatabricksCliTokenSource method using reflection and + * extract the command that would be executed. Note: The CliTokenSource.cmd field is modified by + * OSUtils.getCliExecutableCommand(), so we need to extract it before that transformation. + */ + private List getCliCommand(DatabricksConfig config) throws Exception { + // We'll build the command directly here to test the logic + // This is similar to what getDatabricksCliTokenSource does + String cliPath = config.getDatabricksCliPath(); + if (cliPath == null) { + return null; + } + + java.util.List cmd = + new java.util.ArrayList<>( + java.util.Arrays.asList(cliPath, "auth", "token", "--host", config.getHost())); + if (config.getExperimentalIsUnifiedHost() != null + && config.getExperimentalIsUnifiedHost()) { + cmd.add("--experimental-is-unified-host"); + if (config.getAccountId() != null) { + cmd.add("--account-id"); + cmd.add(config.getAccountId()); + } + if (config.getWorkspaceId() != null) { + cmd.add("--workspace-id"); + cmd.add(config.getWorkspaceId()); + } + } else if (config.getClientType() == ClientType.ACCOUNT) { + cmd.add("--account-id"); + cmd.add(config.getAccountId()); + } + + return cmd; + } + + @Test + void testUnifiedHostWithAccountId() throws Exception { + DatabricksConfig config = + new DatabricksConfig() + .setHost(HOST) + .setAccountId(ACCOUNT_ID) + .setExperimentalIsUnifiedHost(true) + .setDatabricksCliPath("/usr/bin/databricks"); + + List cmd = getCliCommand(config); + + assertNotNull(cmd); + assertTrue(cmd.contains("--experimental-is-unified-host")); + assertTrue(cmd.contains("--account-id")); + assertTrue(cmd.contains(ACCOUNT_ID)); + assertFalse(cmd.contains("--workspace-id")); + } + + @Test + void testUnifiedHostWithWorkspaceId() throws Exception { + DatabricksConfig config = + new DatabricksConfig() + .setHost(HOST) + .setWorkspaceId(WORKSPACE_ID) + .setExperimentalIsUnifiedHost(true) + .setDatabricksCliPath("/usr/bin/databricks"); + + List cmd = getCliCommand(config); + + assertNotNull(cmd); + assertTrue(cmd.contains("--experimental-is-unified-host")); + assertTrue(cmd.contains("--workspace-id")); + assertTrue(cmd.contains(WORKSPACE_ID)); + assertFalse(cmd.contains("--account-id")); + } + + @Test + void testUnifiedHostWithBothIds() throws Exception { + DatabricksConfig config = + new DatabricksConfig() + .setHost(HOST) + .setAccountId(ACCOUNT_ID) + .setWorkspaceId(WORKSPACE_ID) + .setExperimentalIsUnifiedHost(true) + .setDatabricksCliPath("/usr/bin/databricks"); + + List cmd = getCliCommand(config); + + assertNotNull(cmd); + assertTrue(cmd.contains("--experimental-is-unified-host")); + assertTrue(cmd.contains("--account-id")); + assertTrue(cmd.contains(ACCOUNT_ID)); + assertTrue(cmd.contains("--workspace-id")); + assertTrue(cmd.contains(WORKSPACE_ID)); + } + + @Test + void testUnifiedHostWithoutIds() throws Exception { + DatabricksConfig config = + new DatabricksConfig() + .setHost(HOST) + .setExperimentalIsUnifiedHost(true) + .setDatabricksCliPath("/usr/bin/databricks"); + + List cmd = getCliCommand(config); + + assertNotNull(cmd); + assertTrue(cmd.contains("--experimental-is-unified-host")); + assertFalse(cmd.contains("--account-id")); + assertFalse(cmd.contains("--workspace-id")); + } + + @Test + void testAccountClientWithoutUnifiedHost() throws Exception { + DatabricksConfig config = + new DatabricksConfig() + .setHost("https://accounts.cloud.databricks.com") + .setAccountId(ACCOUNT_ID) + .setDatabricksCliPath("/usr/bin/databricks"); + + // ClientType is automatically determined as ACCOUNT for accounts host + assertEquals(ClientType.ACCOUNT, config.getClientType()); + + List cmd = getCliCommand(config); + + assertNotNull(cmd); + assertFalse(cmd.contains("--experimental-is-unified-host")); + assertTrue(cmd.contains("--account-id")); + assertTrue(cmd.contains(ACCOUNT_ID)); + } + + @Test + void testWorkspaceClientWithoutUnifiedHost() throws Exception { + DatabricksConfig config = + new DatabricksConfig().setHost(HOST).setDatabricksCliPath("/usr/bin/databricks"); + + // ClientType is automatically determined as WORKSPACE for workspace host + assertEquals(ClientType.WORKSPACE, config.getClientType()); + + List cmd = getCliCommand(config); + + assertNotNull(cmd); + assertFalse(cmd.contains("--experimental-is-unified-host")); + assertFalse(cmd.contains("--account-id")); + assertFalse(cmd.contains("--workspace-id")); + assertTrue(cmd.contains("auth")); + assertTrue(cmd.contains("token")); + assertTrue(cmd.contains("--host")); + assertTrue(cmd.contains(HOST)); + } + + @Test + void testUnifiedHostFalseUsesRegularBehavior() throws Exception { + DatabricksConfig config = + new DatabricksConfig() + .setHost("https://accounts.cloud.databricks.com") + .setAccountId(ACCOUNT_ID) + .setExperimentalIsUnifiedHost(false) + .setDatabricksCliPath("/usr/bin/databricks"); + + // ClientType is automatically determined as ACCOUNT for accounts host + assertEquals(ClientType.ACCOUNT, config.getClientType()); + + List cmd = getCliCommand(config); + + assertNotNull(cmd); + assertFalse(cmd.contains("--experimental-is-unified-host")); + assertTrue(cmd.contains("--account-id")); + assertTrue(cmd.contains(ACCOUNT_ID)); + } + + @Test + void testConfigureWithNullHost() { + DatabricksCliCredentialsProvider provider = new DatabricksCliCredentialsProvider(); + + DatabricksConfig config = new DatabricksConfig(); + + HeaderFactory header = provider.configure(config); + assertNull(header); + } +} From 49a6f47edc1630480c259b605b6ec6632d9b81de Mon Sep 17 00:00:00 2001 From: Tanmay Rustagi Date: Thu, 22 Jan 2026 15:18:44 +0100 Subject: [PATCH 3/6] - --- .../DatabricksCliCredentialsProvider.java | 2 +- .../DatabricksCliCredentialsProviderTest.java | 164 +++++++++++------- 2 files changed, 104 insertions(+), 62 deletions(-) diff --git a/databricks-sdk-java/src/main/java/com/databricks/sdk/core/DatabricksCliCredentialsProvider.java b/databricks-sdk-java/src/main/java/com/databricks/sdk/core/DatabricksCliCredentialsProvider.java index 57fb22e7c..c25d94277 100644 --- a/databricks-sdk-java/src/main/java/com/databricks/sdk/core/DatabricksCliCredentialsProvider.java +++ b/databricks-sdk-java/src/main/java/com/databricks/sdk/core/DatabricksCliCredentialsProvider.java @@ -20,7 +20,7 @@ public String authType() { return DATABRICKS_CLI; } - private CliTokenSource getDatabricksCliTokenSource(DatabricksConfig config) { + CliTokenSource getDatabricksCliTokenSource(DatabricksConfig config) { String cliPath = config.getDatabricksCliPath(); if (cliPath == null) { cliPath = OSUtils.get(config.getEnv()).getDatabricksCliPath(); diff --git a/databricks-sdk-java/src/test/java/com/databricks/sdk/core/DatabricksCliCredentialsProviderTest.java b/databricks-sdk-java/src/test/java/com/databricks/sdk/core/DatabricksCliCredentialsProviderTest.java index 67de6faf1..2c92da488 100644 --- a/databricks-sdk-java/src/test/java/com/databricks/sdk/core/DatabricksCliCredentialsProviderTest.java +++ b/databricks-sdk-java/src/test/java/com/databricks/sdk/core/DatabricksCliCredentialsProviderTest.java @@ -3,64 +3,89 @@ import static org.junit.jupiter.api.Assertions.*; import com.databricks.sdk.core.oauth.Token; -import com.databricks.sdk.core.utils.Environment; -import java.lang.reflect.Method; import java.time.Instant; +import java.util.ArrayList; +import java.util.Arrays; import java.util.List; import org.junit.jupiter.api.Test; +import org.mockito.Mockito; class DatabricksCliCredentialsProviderTest { - private static final String HOST = "https://db-deco-test.databricks.com"; - private static final String ACCOUNT_ID = "968367da-7edd-44f7-9dea-3e0b20b0ec97"; - private static final String WORKSPACE_ID = "470576644108500"; + private static final String HOST = "https://company.databricks.com"; + private static final String ACCOUNT_ID = "12345678"; + private static final String WORKSPACE_ID = "1234"; + private static final String TOKEN = "test-token"; + private static final String TOKEN_TYPE = "Bearer"; + + private static CliTokenSource mockCliTokenSource() { + CliTokenSource cliTokenSource = Mockito.mock(CliTokenSource.class); + Mockito.when(cliTokenSource.getToken()) + .thenReturn(new Token(TOKEN, TOKEN_TYPE, Instant.now().plusSeconds(3600))); + return cliTokenSource; + } /** - * Helper method to invoke the private getDatabricksCliTokenSource method using reflection and - * extract the command that would be executed. Note: The CliTokenSource.cmd field is modified by - * OSUtils.getCliExecutableCommand(), so we need to extract it before that transformation. + * Test provider that captures the CLI command for verification. This follows the same pattern as + * AzureCliCredentialsProviderTest. Overrides the method to capture the command without actually + * executing the CLI. */ - private List getCliCommand(DatabricksConfig config) throws Exception { - // We'll build the command directly here to test the logic - // This is similar to what getDatabricksCliTokenSource does - String cliPath = config.getDatabricksCliPath(); - if (cliPath == null) { - return null; - } + private static class TestDatabricksCliCredentialsProvider + extends DatabricksCliCredentialsProvider { + private List capturedCommand; + + @Override + CliTokenSource getDatabricksCliTokenSource(DatabricksConfig config) { + String cliPath = config.getDatabricksCliPath(); + if (cliPath == null) { + return null; + } - java.util.List cmd = - new java.util.ArrayList<>( - java.util.Arrays.asList(cliPath, "auth", "token", "--host", config.getHost())); - if (config.getExperimentalIsUnifiedHost() != null - && config.getExperimentalIsUnifiedHost()) { - cmd.add("--experimental-is-unified-host"); - if (config.getAccountId() != null) { + // Build command exactly as the real implementation does (before OSUtils transforms it) + List cmd = + new ArrayList<>(Arrays.asList(cliPath, "auth", "token", "--host", config.getHost())); + if (config.getExperimentalIsUnifiedHost() != null + && config.getExperimentalIsUnifiedHost()) { + cmd.add("--experimental-is-unified-host"); + if (config.getAccountId() != null) { + cmd.add("--account-id"); + cmd.add(config.getAccountId()); + } + if (config.getWorkspaceId() != null) { + cmd.add("--workspace-id"); + cmd.add(config.getWorkspaceId()); + } + } else if (config.getClientType() == ClientType.ACCOUNT) { cmd.add("--account-id"); cmd.add(config.getAccountId()); } - if (config.getWorkspaceId() != null) { - cmd.add("--workspace-id"); - cmd.add(config.getWorkspaceId()); - } - } else if (config.getClientType() == ClientType.ACCOUNT) { - cmd.add("--account-id"); - cmd.add(config.getAccountId()); + + this.capturedCommand = cmd; + + // Return mock token source instead of creating real one + return mockCliTokenSource(); } - return cmd; + public List getCapturedCommand() { + return capturedCommand; + } } @Test - void testUnifiedHostWithAccountId() throws Exception { + void testUnifiedHostWithAccountId() { + TestDatabricksCliCredentialsProvider provider = new TestDatabricksCliCredentialsProvider(); DatabricksConfig config = new DatabricksConfig() .setHost(HOST) .setAccountId(ACCOUNT_ID) .setExperimentalIsUnifiedHost(true) - .setDatabricksCliPath("/usr/bin/databricks"); + .setDatabricksCliPath("/usr/bin/databricks") + .setCredentialsProvider(provider); - List cmd = getCliCommand(config); + HeaderFactory header = provider.configure(config); + assertNotNull(header); + List cmd = provider.getCapturedCommand(); assertNotNull(cmd); assertTrue(cmd.contains("--experimental-is-unified-host")); assertTrue(cmd.contains("--account-id")); @@ -69,16 +94,20 @@ void testUnifiedHostWithAccountId() throws Exception { } @Test - void testUnifiedHostWithWorkspaceId() throws Exception { + void testUnifiedHostWithWorkspaceId() { + TestDatabricksCliCredentialsProvider provider = new TestDatabricksCliCredentialsProvider(); DatabricksConfig config = new DatabricksConfig() .setHost(HOST) .setWorkspaceId(WORKSPACE_ID) .setExperimentalIsUnifiedHost(true) - .setDatabricksCliPath("/usr/bin/databricks"); + .setDatabricksCliPath("/usr/bin/databricks") + .setCredentialsProvider(provider); - List cmd = getCliCommand(config); + HeaderFactory header = provider.configure(config); + assertNotNull(header); + List cmd = provider.getCapturedCommand(); assertNotNull(cmd); assertTrue(cmd.contains("--experimental-is-unified-host")); assertTrue(cmd.contains("--workspace-id")); @@ -87,17 +116,21 @@ void testUnifiedHostWithWorkspaceId() throws Exception { } @Test - void testUnifiedHostWithBothIds() throws Exception { + void testUnifiedHostWithBothIds() { + TestDatabricksCliCredentialsProvider provider = new TestDatabricksCliCredentialsProvider(); DatabricksConfig config = new DatabricksConfig() .setHost(HOST) .setAccountId(ACCOUNT_ID) .setWorkspaceId(WORKSPACE_ID) .setExperimentalIsUnifiedHost(true) - .setDatabricksCliPath("/usr/bin/databricks"); + .setDatabricksCliPath("/usr/bin/databricks") + .setCredentialsProvider(provider); - List cmd = getCliCommand(config); + HeaderFactory header = provider.configure(config); + assertNotNull(header); + List cmd = provider.getCapturedCommand(); assertNotNull(cmd); assertTrue(cmd.contains("--experimental-is-unified-host")); assertTrue(cmd.contains("--account-id")); @@ -107,15 +140,19 @@ void testUnifiedHostWithBothIds() throws Exception { } @Test - void testUnifiedHostWithoutIds() throws Exception { + void testUnifiedHostWithoutIds() { + TestDatabricksCliCredentialsProvider provider = new TestDatabricksCliCredentialsProvider(); DatabricksConfig config = new DatabricksConfig() .setHost(HOST) .setExperimentalIsUnifiedHost(true) - .setDatabricksCliPath("/usr/bin/databricks"); + .setDatabricksCliPath("/usr/bin/databricks") + .setCredentialsProvider(provider); - List cmd = getCliCommand(config); + HeaderFactory header = provider.configure(config); + assertNotNull(header); + List cmd = provider.getCapturedCommand(); assertNotNull(cmd); assertTrue(cmd.contains("--experimental-is-unified-host")); assertFalse(cmd.contains("--account-id")); @@ -123,18 +160,19 @@ void testUnifiedHostWithoutIds() throws Exception { } @Test - void testAccountClientWithoutUnifiedHost() throws Exception { + void testAccountClientWithoutUnifiedHost() { + TestDatabricksCliCredentialsProvider provider = new TestDatabricksCliCredentialsProvider(); DatabricksConfig config = new DatabricksConfig() .setHost("https://accounts.cloud.databricks.com") .setAccountId(ACCOUNT_ID) - .setDatabricksCliPath("/usr/bin/databricks"); + .setDatabricksCliPath("/usr/bin/databricks") + .setCredentialsProvider(provider); - // ClientType is automatically determined as ACCOUNT for accounts host - assertEquals(ClientType.ACCOUNT, config.getClientType()); - - List cmd = getCliCommand(config); + HeaderFactory header = provider.configure(config); + assertNotNull(header); + List cmd = provider.getCapturedCommand(); assertNotNull(cmd); assertFalse(cmd.contains("--experimental-is-unified-host")); assertTrue(cmd.contains("--account-id")); @@ -142,15 +180,18 @@ void testAccountClientWithoutUnifiedHost() throws Exception { } @Test - void testWorkspaceClientWithoutUnifiedHost() throws Exception { + void testWorkspaceClientWithoutUnifiedHost() { + TestDatabricksCliCredentialsProvider provider = new TestDatabricksCliCredentialsProvider(); DatabricksConfig config = - new DatabricksConfig().setHost(HOST).setDatabricksCliPath("/usr/bin/databricks"); - - // ClientType is automatically determined as WORKSPACE for workspace host - assertEquals(ClientType.WORKSPACE, config.getClientType()); + new DatabricksConfig() + .setHost(HOST) + .setDatabricksCliPath("/usr/bin/databricks") + .setCredentialsProvider(provider); - List cmd = getCliCommand(config); + HeaderFactory header = provider.configure(config); + assertNotNull(header); + List cmd = provider.getCapturedCommand(); assertNotNull(cmd); assertFalse(cmd.contains("--experimental-is-unified-host")); assertFalse(cmd.contains("--account-id")); @@ -162,19 +203,20 @@ void testWorkspaceClientWithoutUnifiedHost() throws Exception { } @Test - void testUnifiedHostFalseUsesRegularBehavior() throws Exception { + void testUnifiedHostFalseUsesRegularBehavior() { + TestDatabricksCliCredentialsProvider provider = new TestDatabricksCliCredentialsProvider(); DatabricksConfig config = new DatabricksConfig() .setHost("https://accounts.cloud.databricks.com") .setAccountId(ACCOUNT_ID) .setExperimentalIsUnifiedHost(false) - .setDatabricksCliPath("/usr/bin/databricks"); - - // ClientType is automatically determined as ACCOUNT for accounts host - assertEquals(ClientType.ACCOUNT, config.getClientType()); + .setDatabricksCliPath("/usr/bin/databricks") + .setCredentialsProvider(provider); - List cmd = getCliCommand(config); + HeaderFactory header = provider.configure(config); + assertNotNull(header); + List cmd = provider.getCapturedCommand(); assertNotNull(cmd); assertFalse(cmd.contains("--experimental-is-unified-host")); assertTrue(cmd.contains("--account-id")); @@ -184,10 +226,10 @@ void testUnifiedHostFalseUsesRegularBehavior() throws Exception { @Test void testConfigureWithNullHost() { DatabricksCliCredentialsProvider provider = new DatabricksCliCredentialsProvider(); - DatabricksConfig config = new DatabricksConfig(); HeaderFactory header = provider.configure(config); + assertNull(header); } } From 7b32de5ed205465c188c7334616d8116eb1bfb67 Mon Sep 17 00:00:00 2001 From: Tanmay Rustagi Date: Thu, 22 Jan 2026 15:40:17 +0100 Subject: [PATCH 4/6] - --- .../DatabricksCliCredentialsProvider.java | 30 ++- .../DatabricksCliCredentialsProviderTest.java | 235 ------------------ 2 files changed, 21 insertions(+), 244 deletions(-) delete mode 100644 databricks-sdk-java/src/test/java/com/databricks/sdk/core/DatabricksCliCredentialsProviderTest.java diff --git a/databricks-sdk-java/src/main/java/com/databricks/sdk/core/DatabricksCliCredentialsProvider.java b/databricks-sdk-java/src/main/java/com/databricks/sdk/core/DatabricksCliCredentialsProvider.java index c25d94277..7fc505583 100644 --- a/databricks-sdk-java/src/main/java/com/databricks/sdk/core/DatabricksCliCredentialsProvider.java +++ b/databricks-sdk-java/src/main/java/com/databricks/sdk/core/DatabricksCliCredentialsProvider.java @@ -20,15 +20,14 @@ public String authType() { return DATABRICKS_CLI; } - CliTokenSource getDatabricksCliTokenSource(DatabricksConfig config) { - String cliPath = config.getDatabricksCliPath(); - if (cliPath == null) { - cliPath = OSUtils.get(config.getEnv()).getDatabricksCliPath(); - } - if (cliPath == null) { - LOG.debug("Databricks CLI could not be found"); - return null; - } + /** + * Builds the CLI command arguments for the databricks auth token command. + * + * @param cliPath Path to the databricks CLI executable + * @param config Configuration containing host, account ID, workspace ID, etc. + * @return List of command arguments + */ + List buildCliCommand(String cliPath, DatabricksConfig config) { List cmd = new ArrayList<>(Arrays.asList(cliPath, "auth", "token", "--host", config.getHost())); if (config.getExperimentalIsUnifiedHost() != null && config.getExperimentalIsUnifiedHost()) { @@ -46,6 +45,19 @@ CliTokenSource getDatabricksCliTokenSource(DatabricksConfig config) { cmd.add("--account-id"); cmd.add(config.getAccountId()); } + return cmd; + } + + private CliTokenSource getDatabricksCliTokenSource(DatabricksConfig config) { + String cliPath = config.getDatabricksCliPath(); + if (cliPath == null) { + cliPath = OSUtils.get(config.getEnv()).getDatabricksCliPath(); + } + if (cliPath == null) { + LOG.debug("Databricks CLI could not be found"); + return null; + } + List cmd = buildCliCommand(cliPath, config); return new CliTokenSource(cmd, "token_type", "access_token", "expiry", config.getEnv()); } diff --git a/databricks-sdk-java/src/test/java/com/databricks/sdk/core/DatabricksCliCredentialsProviderTest.java b/databricks-sdk-java/src/test/java/com/databricks/sdk/core/DatabricksCliCredentialsProviderTest.java deleted file mode 100644 index 2c92da488..000000000 --- a/databricks-sdk-java/src/test/java/com/databricks/sdk/core/DatabricksCliCredentialsProviderTest.java +++ /dev/null @@ -1,235 +0,0 @@ -package com.databricks.sdk.core; - -import static org.junit.jupiter.api.Assertions.*; - -import com.databricks.sdk.core.oauth.Token; -import java.time.Instant; -import java.util.ArrayList; -import java.util.Arrays; -import java.util.List; -import org.junit.jupiter.api.Test; -import org.mockito.Mockito; - -class DatabricksCliCredentialsProviderTest { - - private static final String HOST = "https://company.databricks.com"; - private static final String ACCOUNT_ID = "12345678"; - private static final String WORKSPACE_ID = "1234"; - private static final String TOKEN = "test-token"; - private static final String TOKEN_TYPE = "Bearer"; - - private static CliTokenSource mockCliTokenSource() { - CliTokenSource cliTokenSource = Mockito.mock(CliTokenSource.class); - Mockito.when(cliTokenSource.getToken()) - .thenReturn(new Token(TOKEN, TOKEN_TYPE, Instant.now().plusSeconds(3600))); - return cliTokenSource; - } - - /** - * Test provider that captures the CLI command for verification. This follows the same pattern as - * AzureCliCredentialsProviderTest. Overrides the method to capture the command without actually - * executing the CLI. - */ - private static class TestDatabricksCliCredentialsProvider - extends DatabricksCliCredentialsProvider { - private List capturedCommand; - - @Override - CliTokenSource getDatabricksCliTokenSource(DatabricksConfig config) { - String cliPath = config.getDatabricksCliPath(); - if (cliPath == null) { - return null; - } - - // Build command exactly as the real implementation does (before OSUtils transforms it) - List cmd = - new ArrayList<>(Arrays.asList(cliPath, "auth", "token", "--host", config.getHost())); - if (config.getExperimentalIsUnifiedHost() != null - && config.getExperimentalIsUnifiedHost()) { - cmd.add("--experimental-is-unified-host"); - if (config.getAccountId() != null) { - cmd.add("--account-id"); - cmd.add(config.getAccountId()); - } - if (config.getWorkspaceId() != null) { - cmd.add("--workspace-id"); - cmd.add(config.getWorkspaceId()); - } - } else if (config.getClientType() == ClientType.ACCOUNT) { - cmd.add("--account-id"); - cmd.add(config.getAccountId()); - } - - this.capturedCommand = cmd; - - // Return mock token source instead of creating real one - return mockCliTokenSource(); - } - - public List getCapturedCommand() { - return capturedCommand; - } - } - - @Test - void testUnifiedHostWithAccountId() { - TestDatabricksCliCredentialsProvider provider = new TestDatabricksCliCredentialsProvider(); - DatabricksConfig config = - new DatabricksConfig() - .setHost(HOST) - .setAccountId(ACCOUNT_ID) - .setExperimentalIsUnifiedHost(true) - .setDatabricksCliPath("/usr/bin/databricks") - .setCredentialsProvider(provider); - - HeaderFactory header = provider.configure(config); - - assertNotNull(header); - List cmd = provider.getCapturedCommand(); - assertNotNull(cmd); - assertTrue(cmd.contains("--experimental-is-unified-host")); - assertTrue(cmd.contains("--account-id")); - assertTrue(cmd.contains(ACCOUNT_ID)); - assertFalse(cmd.contains("--workspace-id")); - } - - @Test - void testUnifiedHostWithWorkspaceId() { - TestDatabricksCliCredentialsProvider provider = new TestDatabricksCliCredentialsProvider(); - DatabricksConfig config = - new DatabricksConfig() - .setHost(HOST) - .setWorkspaceId(WORKSPACE_ID) - .setExperimentalIsUnifiedHost(true) - .setDatabricksCliPath("/usr/bin/databricks") - .setCredentialsProvider(provider); - - HeaderFactory header = provider.configure(config); - - assertNotNull(header); - List cmd = provider.getCapturedCommand(); - assertNotNull(cmd); - assertTrue(cmd.contains("--experimental-is-unified-host")); - assertTrue(cmd.contains("--workspace-id")); - assertTrue(cmd.contains(WORKSPACE_ID)); - assertFalse(cmd.contains("--account-id")); - } - - @Test - void testUnifiedHostWithBothIds() { - TestDatabricksCliCredentialsProvider provider = new TestDatabricksCliCredentialsProvider(); - DatabricksConfig config = - new DatabricksConfig() - .setHost(HOST) - .setAccountId(ACCOUNT_ID) - .setWorkspaceId(WORKSPACE_ID) - .setExperimentalIsUnifiedHost(true) - .setDatabricksCliPath("/usr/bin/databricks") - .setCredentialsProvider(provider); - - HeaderFactory header = provider.configure(config); - - assertNotNull(header); - List cmd = provider.getCapturedCommand(); - assertNotNull(cmd); - assertTrue(cmd.contains("--experimental-is-unified-host")); - assertTrue(cmd.contains("--account-id")); - assertTrue(cmd.contains(ACCOUNT_ID)); - assertTrue(cmd.contains("--workspace-id")); - assertTrue(cmd.contains(WORKSPACE_ID)); - } - - @Test - void testUnifiedHostWithoutIds() { - TestDatabricksCliCredentialsProvider provider = new TestDatabricksCliCredentialsProvider(); - DatabricksConfig config = - new DatabricksConfig() - .setHost(HOST) - .setExperimentalIsUnifiedHost(true) - .setDatabricksCliPath("/usr/bin/databricks") - .setCredentialsProvider(provider); - - HeaderFactory header = provider.configure(config); - - assertNotNull(header); - List cmd = provider.getCapturedCommand(); - assertNotNull(cmd); - assertTrue(cmd.contains("--experimental-is-unified-host")); - assertFalse(cmd.contains("--account-id")); - assertFalse(cmd.contains("--workspace-id")); - } - - @Test - void testAccountClientWithoutUnifiedHost() { - TestDatabricksCliCredentialsProvider provider = new TestDatabricksCliCredentialsProvider(); - DatabricksConfig config = - new DatabricksConfig() - .setHost("https://accounts.cloud.databricks.com") - .setAccountId(ACCOUNT_ID) - .setDatabricksCliPath("/usr/bin/databricks") - .setCredentialsProvider(provider); - - HeaderFactory header = provider.configure(config); - - assertNotNull(header); - List cmd = provider.getCapturedCommand(); - assertNotNull(cmd); - assertFalse(cmd.contains("--experimental-is-unified-host")); - assertTrue(cmd.contains("--account-id")); - assertTrue(cmd.contains(ACCOUNT_ID)); - } - - @Test - void testWorkspaceClientWithoutUnifiedHost() { - TestDatabricksCliCredentialsProvider provider = new TestDatabricksCliCredentialsProvider(); - DatabricksConfig config = - new DatabricksConfig() - .setHost(HOST) - .setDatabricksCliPath("/usr/bin/databricks") - .setCredentialsProvider(provider); - - HeaderFactory header = provider.configure(config); - - assertNotNull(header); - List cmd = provider.getCapturedCommand(); - assertNotNull(cmd); - assertFalse(cmd.contains("--experimental-is-unified-host")); - assertFalse(cmd.contains("--account-id")); - assertFalse(cmd.contains("--workspace-id")); - assertTrue(cmd.contains("auth")); - assertTrue(cmd.contains("token")); - assertTrue(cmd.contains("--host")); - assertTrue(cmd.contains(HOST)); - } - - @Test - void testUnifiedHostFalseUsesRegularBehavior() { - TestDatabricksCliCredentialsProvider provider = new TestDatabricksCliCredentialsProvider(); - DatabricksConfig config = - new DatabricksConfig() - .setHost("https://accounts.cloud.databricks.com") - .setAccountId(ACCOUNT_ID) - .setExperimentalIsUnifiedHost(false) - .setDatabricksCliPath("/usr/bin/databricks") - .setCredentialsProvider(provider); - - HeaderFactory header = provider.configure(config); - - assertNotNull(header); - List cmd = provider.getCapturedCommand(); - assertNotNull(cmd); - assertFalse(cmd.contains("--experimental-is-unified-host")); - assertTrue(cmd.contains("--account-id")); - assertTrue(cmd.contains(ACCOUNT_ID)); - } - - @Test - void testConfigureWithNullHost() { - DatabricksCliCredentialsProvider provider = new DatabricksCliCredentialsProvider(); - DatabricksConfig config = new DatabricksConfig(); - - HeaderFactory header = provider.configure(config); - - assertNull(header); - } -} From 4d1f8b30fd88ab5722a08683c3baa3df8a63380c Mon Sep 17 00:00:00 2001 From: Tanmay Rustagi Date: Thu, 22 Jan 2026 15:44:41 +0100 Subject: [PATCH 5/6] - --- .../DatabricksCliCredentialsProviderTest.java | 140 ++++++++++++++++++ 1 file changed, 140 insertions(+) create mode 100644 databricks-sdk-java/src/test/java/com/databricks/sdk/core/DatabricksCliCredentialsProviderTest.java diff --git a/databricks-sdk-java/src/test/java/com/databricks/sdk/core/DatabricksCliCredentialsProviderTest.java b/databricks-sdk-java/src/test/java/com/databricks/sdk/core/DatabricksCliCredentialsProviderTest.java new file mode 100644 index 000000000..6bd141eec --- /dev/null +++ b/databricks-sdk-java/src/test/java/com/databricks/sdk/core/DatabricksCliCredentialsProviderTest.java @@ -0,0 +1,140 @@ +package com.databricks.sdk.core; + +import static org.junit.jupiter.api.Assertions.*; + +import java.util.Arrays; +import java.util.List; +import org.junit.jupiter.api.Test; + +class DatabricksCliCredentialsProviderTest { + + private static final String CLI_PATH = "/usr/local/bin/databricks"; + private static final String HOST = "https://my-workspace.cloud.databricks.com"; + private static final String ACCOUNT_HOST = "https://accounts.cloud.databricks.com"; + private static final String UNIFIED_HOST = "https://unified.databricks.com"; + private static final String ACCOUNT_ID = "test-account-123"; + private static final String WORKSPACE_ID = "987654321"; + + private final DatabricksCliCredentialsProvider provider = new DatabricksCliCredentialsProvider(); + + @Test + void testBuildCliCommand_WorkspaceHost() { + DatabricksConfig config = new DatabricksConfig().setHost(HOST); + + List cmd = provider.buildCliCommand(CLI_PATH, config); + + assertEquals(Arrays.asList(CLI_PATH, "auth", "token", "--host", HOST), cmd); + } + + @Test + void testBuildCliCommand_AccountHost() { + DatabricksConfig config = new DatabricksConfig().setHost(ACCOUNT_HOST).setAccountId(ACCOUNT_ID); + + List cmd = provider.buildCliCommand(CLI_PATH, config); + + assertEquals( + Arrays.asList(CLI_PATH, "auth", "token", "--host", ACCOUNT_HOST, "--account-id", ACCOUNT_ID), + cmd); + } + + @Test + void testBuildCliCommand_UnifiedHost_WithAccountIdAndWorkspaceId() { + DatabricksConfig config = + new DatabricksConfig() + .setHost(UNIFIED_HOST) + .setExperimentalIsUnifiedHost(true) + .setAccountId(ACCOUNT_ID) + .setWorkspaceId(WORKSPACE_ID); + + List cmd = provider.buildCliCommand(CLI_PATH, config); + + assertEquals( + Arrays.asList( + CLI_PATH, + "auth", + "token", + "--host", + UNIFIED_HOST, + "--experimental-is-unified-host", + "--account-id", + ACCOUNT_ID, + "--workspace-id", + WORKSPACE_ID), + cmd); + } + + @Test + void testBuildCliCommand_UnifiedHost_WithAccountIdOnly() { + DatabricksConfig config = + new DatabricksConfig() + .setHost(UNIFIED_HOST) + .setExperimentalIsUnifiedHost(true) + .setAccountId(ACCOUNT_ID); + + List cmd = provider.buildCliCommand(CLI_PATH, config); + + assertEquals( + Arrays.asList( + CLI_PATH, + "auth", + "token", + "--host", + UNIFIED_HOST, + "--experimental-is-unified-host", + "--account-id", + ACCOUNT_ID), + cmd); + } + + @Test + void testBuildCliCommand_UnifiedHost_WithWorkspaceIdOnly() { + DatabricksConfig config = + new DatabricksConfig() + .setHost(UNIFIED_HOST) + .setExperimentalIsUnifiedHost(true) + .setWorkspaceId(WORKSPACE_ID); + + List cmd = provider.buildCliCommand(CLI_PATH, config); + + assertEquals( + Arrays.asList( + CLI_PATH, + "auth", + "token", + "--host", + UNIFIED_HOST, + "--experimental-is-unified-host", + "--workspace-id", + WORKSPACE_ID), + cmd); + } + + @Test + void testBuildCliCommand_UnifiedHost_WithNoAccountIdOrWorkspaceId() { + DatabricksConfig config = + new DatabricksConfig().setHost(UNIFIED_HOST).setExperimentalIsUnifiedHost(true); + + List cmd = provider.buildCliCommand(CLI_PATH, config); + + assertEquals( + Arrays.asList( + CLI_PATH, "auth", "token", "--host", UNIFIED_HOST, "--experimental-is-unified-host"), + cmd); + } + + @Test + void testBuildCliCommand_UnifiedHostFalse_WithAccountHost() { + // When experimentalIsUnifiedHost is explicitly false, should fall back to account-id logic + DatabricksConfig config = + new DatabricksConfig() + .setHost(ACCOUNT_HOST) + .setExperimentalIsUnifiedHost(false) + .setAccountId(ACCOUNT_ID); + + List cmd = provider.buildCliCommand(CLI_PATH, config); + + assertEquals( + Arrays.asList(CLI_PATH, "auth", "token", "--host", ACCOUNT_HOST, "--account-id", ACCOUNT_ID), + cmd); + } +} From 2b727c741fb970252a4cf7cd049ba3d3c4ccd24a Mon Sep 17 00:00:00 2001 From: Tanmay Rustagi Date: Thu, 22 Jan 2026 15:46:09 +0100 Subject: [PATCH 6/6] fmt --- .../sdk/core/DatabricksCliCredentialsProviderTest.java | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) diff --git a/databricks-sdk-java/src/test/java/com/databricks/sdk/core/DatabricksCliCredentialsProviderTest.java b/databricks-sdk-java/src/test/java/com/databricks/sdk/core/DatabricksCliCredentialsProviderTest.java index 6bd141eec..478948d82 100644 --- a/databricks-sdk-java/src/test/java/com/databricks/sdk/core/DatabricksCliCredentialsProviderTest.java +++ b/databricks-sdk-java/src/test/java/com/databricks/sdk/core/DatabricksCliCredentialsProviderTest.java @@ -33,7 +33,8 @@ void testBuildCliCommand_AccountHost() { List cmd = provider.buildCliCommand(CLI_PATH, config); assertEquals( - Arrays.asList(CLI_PATH, "auth", "token", "--host", ACCOUNT_HOST, "--account-id", ACCOUNT_ID), + Arrays.asList( + CLI_PATH, "auth", "token", "--host", ACCOUNT_HOST, "--account-id", ACCOUNT_ID), cmd); } @@ -134,7 +135,8 @@ void testBuildCliCommand_UnifiedHostFalse_WithAccountHost() { List cmd = provider.buildCliCommand(CLI_PATH, config); assertEquals( - Arrays.asList(CLI_PATH, "auth", "token", "--host", ACCOUNT_HOST, "--account-id", ACCOUNT_ID), + Arrays.asList( + CLI_PATH, "auth", "token", "--host", ACCOUNT_HOST, "--account-id", ACCOUNT_ID), cmd); } }