Skip to content

Commit 3574190

Browse files
tanmay-dbtanmay-db
authored
[Feature] Support CLI Profiles for unified host (#629)
## What changes are proposed in this pull request? - Add support to use profiles and tokens generated from the databricks cli auth login command. ## How is this tested? - Unit tests - Manually tested using CLI generated profile NO_CHANGELOG=true
1 parent f6fb49c commit 3574190

File tree

2 files changed

+171
-6
lines changed

2 files changed

+171
-6
lines changed

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

Lines changed: 29 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -20,6 +20,34 @@ public String authType() {
2020
return DATABRICKS_CLI;
2121
}
2222

23+
/**
24+
* Builds the CLI command arguments for the databricks auth token command.
25+
*
26+
* @param cliPath Path to the databricks CLI executable
27+
* @param config Configuration containing host, account ID, workspace ID, etc.
28+
* @return List of command arguments
29+
*/
30+
List<String> buildCliCommand(String cliPath, DatabricksConfig config) {
31+
List<String> cmd =
32+
new ArrayList<>(Arrays.asList(cliPath, "auth", "token", "--host", config.getHost()));
33+
if (config.getExperimentalIsUnifiedHost() != null && config.getExperimentalIsUnifiedHost()) {
34+
// For unified hosts, pass account_id, workspace_id, and experimental flag
35+
cmd.add("--experimental-is-unified-host");
36+
if (config.getAccountId() != null) {
37+
cmd.add("--account-id");
38+
cmd.add(config.getAccountId());
39+
}
40+
if (config.getWorkspaceId() != null) {
41+
cmd.add("--workspace-id");
42+
cmd.add(config.getWorkspaceId());
43+
}
44+
} else if (config.getClientType() == ClientType.ACCOUNT) {
45+
cmd.add("--account-id");
46+
cmd.add(config.getAccountId());
47+
}
48+
return cmd;
49+
}
50+
2351
private CliTokenSource getDatabricksCliTokenSource(DatabricksConfig config) {
2452
String cliPath = config.getDatabricksCliPath();
2553
if (cliPath == null) {
@@ -29,12 +57,7 @@ private CliTokenSource getDatabricksCliTokenSource(DatabricksConfig config) {
2957
LOG.debug("Databricks CLI could not be found");
3058
return null;
3159
}
32-
List<String> cmd =
33-
new ArrayList<>(Arrays.asList(cliPath, "auth", "token", "--host", config.getHost()));
34-
if (config.getClientType() == ClientType.ACCOUNT) {
35-
cmd.add("--account-id");
36-
cmd.add(config.getAccountId());
37-
}
60+
List<String> cmd = buildCliCommand(cliPath, config);
3861
return new CliTokenSource(cmd, "token_type", "access_token", "expiry", config.getEnv());
3962
}
4063

‎databricks-sdk-java/src/test/java/com/databricks/sdk/core/DatabricksCliCredentialsProviderTest.java‎

Lines changed: 142 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,142 @@
1+
package com.databricks.sdk.core;
2+
3+
import static org.junit.jupiter.api.Assertions.*;
4+
5+
import java.util.Arrays;
6+
import java.util.List;
7+
import org.junit.jupiter.api.Test;
8+
9+
class DatabricksCliCredentialsProviderTest {
10+
11+
private static final String CLI_PATH = "/usr/local/bin/databricks";
12+
private static final String HOST = "https://my-workspace.cloud.databricks.com";
13+
private static final String ACCOUNT_HOST = "https://accounts.cloud.databricks.com";
14+
private static final String UNIFIED_HOST = "https://unified.databricks.com";
15+
private static final String ACCOUNT_ID = "test-account-123";
16+
private static final String WORKSPACE_ID = "987654321";
17+
18+
private final DatabricksCliCredentialsProvider provider = new DatabricksCliCredentialsProvider();
19+
20+
@Test
21+
void testBuildCliCommand_WorkspaceHost() {
22+
DatabricksConfig config = new DatabricksConfig().setHost(HOST);
23+
24+
List<String> cmd = provider.buildCliCommand(CLI_PATH, config);
25+
26+
assertEquals(Arrays.asList(CLI_PATH, "auth", "token", "--host", HOST), cmd);
27+
}
28+
29+
@Test
30+
void testBuildCliCommand_AccountHost() {
31+
DatabricksConfig config = new DatabricksConfig().setHost(ACCOUNT_HOST).setAccountId(ACCOUNT_ID);
32+
33+
List<String> cmd = provider.buildCliCommand(CLI_PATH, config);
34+
35+
assertEquals(
36+
Arrays.asList(
37+
CLI_PATH, "auth", "token", "--host", ACCOUNT_HOST, "--account-id", ACCOUNT_ID),
38+
cmd);
39+
}
40+
41+
@Test
42+
void testBuildCliCommand_UnifiedHost_WithAccountIdAndWorkspaceId() {
43+
DatabricksConfig config =
44+
new DatabricksConfig()
45+
.setHost(UNIFIED_HOST)
46+
.setExperimentalIsUnifiedHost(true)
47+
.setAccountId(ACCOUNT_ID)
48+
.setWorkspaceId(WORKSPACE_ID);
49+
50+
List<String> cmd = provider.buildCliCommand(CLI_PATH, config);
51+
52+
assertEquals(
53+
Arrays.asList(
54+
CLI_PATH,
55+
"auth",
56+
"token",
57+
"--host",
58+
UNIFIED_HOST,
59+
"--experimental-is-unified-host",
60+
"--account-id",
61+
ACCOUNT_ID,
62+
"--workspace-id",
63+
WORKSPACE_ID),
64+
cmd);
65+
}
66+
67+
@Test
68+
void testBuildCliCommand_UnifiedHost_WithAccountIdOnly() {
69+
DatabricksConfig config =
70+
new DatabricksConfig()
71+
.setHost(UNIFIED_HOST)
72+
.setExperimentalIsUnifiedHost(true)
73+
.setAccountId(ACCOUNT_ID);
74+
75+
List<String> cmd = provider.buildCliCommand(CLI_PATH, config);
76+
77+
assertEquals(
78+
Arrays.asList(
79+
CLI_PATH,
80+
"auth",
81+
"token",
82+
"--host",
83+
UNIFIED_HOST,
84+
"--experimental-is-unified-host",
85+
"--account-id",
86+
ACCOUNT_ID),
87+
cmd);
88+
}
89+
90+
@Test
91+
void testBuildCliCommand_UnifiedHost_WithWorkspaceIdOnly() {
92+
DatabricksConfig config =
93+
new DatabricksConfig()
94+
.setHost(UNIFIED_HOST)
95+
.setExperimentalIsUnifiedHost(true)
96+
.setWorkspaceId(WORKSPACE_ID);
97+
98+
List<String> cmd = provider.buildCliCommand(CLI_PATH, config);
99+
100+
assertEquals(
101+
Arrays.asList(
102+
CLI_PATH,
103+
"auth",
104+
"token",
105+
"--host",
106+
UNIFIED_HOST,
107+
"--experimental-is-unified-host",
108+
"--workspace-id",
109+
WORKSPACE_ID),
110+
cmd);
111+
}
112+
113+
@Test
114+
void testBuildCliCommand_UnifiedHost_WithNoAccountIdOrWorkspaceId() {
115+
DatabricksConfig config =
116+
new DatabricksConfig().setHost(UNIFIED_HOST).setExperimentalIsUnifiedHost(true);
117+
118+
List<String> cmd = provider.buildCliCommand(CLI_PATH, config);
119+
120+
assertEquals(
121+
Arrays.asList(
122+
CLI_PATH, "auth", "token", "--host", UNIFIED_HOST, "--experimental-is-unified-host"),
123+
cmd);
124+
}
125+
126+
@Test
127+
void testBuildCliCommand_UnifiedHostFalse_WithAccountHost() {
128+
// When experimentalIsUnifiedHost is explicitly false, should fall back to account-id logic
129+
DatabricksConfig config =
130+
new DatabricksConfig()
131+
.setHost(ACCOUNT_HOST)
132+
.setExperimentalIsUnifiedHost(false)
133+
.setAccountId(ACCOUNT_ID);
134+
135+
List<String> cmd = provider.buildCliCommand(CLI_PATH, config);
136+
137+
assertEquals(
138+
Arrays.asList(
139+
CLI_PATH, "auth", "token", "--host", ACCOUNT_HOST, "--account-id", ACCOUNT_ID),
140+
cmd);
141+
}
142+
}

0 commit comments

Comments
 (0)