Skip to content

Commit d98ea37

Browse files
hectorcast-dbhectorcast-db
committed
wip
1 parent 7f41821 commit d98ea37

File tree

3 files changed

+137
-25
lines changed

3 files changed

+137
-25
lines changed

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

Lines changed: 70 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -18,27 +18,64 @@
1818
import java.util.Arrays;
1919
import java.util.List;
2020
import org.apache.commons.io.IOUtils;
21+
import org.slf4j.Logger;
22+
import org.slf4j.LoggerFactory;
2123

2224
@InternalApi
2325
public class CliTokenSource implements TokenSource {
26+
private static final Logger LOG = LoggerFactory.getLogger(CliTokenSource.class);
27+
2428
private List<String> cmd;
2529
private String tokenTypeField;
2630
private String accessTokenField;
2731
private String expiryField;
2832
private Environment env;
33+
// fallbackCmd is tried when the primary command fails with "unknown flag: --profile",
34+
// indicating the CLI is too old to support --profile. Can be removed once support
35+
// for CLI versions predating --profile is dropped.
36+
// See: https://github.com/databricks/databricks-sdk-go/pull/1497
37+
private List<String> fallbackCmd;
38+
39+
/**
40+
* Internal exception that carries the clean stderr message but exposes full output for checks.
41+
*/
42+
static class CliCommandException extends IOException {
43+
private final String fullOutput;
44+
45+
CliCommandException(String message, String fullOutput) {
46+
super(message);
47+
this.fullOutput = fullOutput;
48+
}
49+
50+
String getFullOutput() {
51+
return fullOutput;
52+
}
53+
}
2954

3055
public CliTokenSource(
3156
List<String> cmd,
3257
String tokenTypeField,
3358
String accessTokenField,
3459
String expiryField,
3560
Environment env) {
61+
this(cmd, tokenTypeField, accessTokenField, expiryField, env, null);
62+
}
63+
64+
public CliTokenSource(
65+
List<String> cmd,
66+
String tokenTypeField,
67+
String accessTokenField,
68+
String expiryField,
69+
Environment env,
70+
List<String> fallbackCmd) {
3671
super();
3772
this.cmd = OSUtils.get(env).getCliExecutableCommand(cmd);
3873
this.tokenTypeField = tokenTypeField;
3974
this.accessTokenField = accessTokenField;
4075
this.expiryField = expiryField;
4176
this.env = env;
77+
this.fallbackCmd =
78+
fallbackCmd != null ? OSUtils.get(env).getCliExecutableCommand(fallbackCmd) : null;
4279
}
4380

4481
/**
@@ -87,10 +124,9 @@ private String getProcessStream(InputStream stream) throws IOException {
87124
return new String(bytes);
88125
}
89126

90-
@Override
91-
public Token getToken() {
127+
private Token execCliCommand(List<String> cmdToRun) throws IOException {
92128
try {
93-
ProcessBuilder processBuilder = new ProcessBuilder(cmd);
129+
ProcessBuilder processBuilder = new ProcessBuilder(cmdToRun);
94130
processBuilder.environment().putAll(env.getEnv());
95131
Process process = processBuilder.start();
96132
String stdout = getProcessStream(process.getInputStream());
@@ -99,9 +135,11 @@ public Token getToken() {
99135
if (exitCode != 0) {
100136
if (stderr.contains("not found")) {
101137
throw new DatabricksException(stderr);
102-
} else {
103-
throw new IOException(stderr);
104138
}
139+
// getMessage() returns the clean stderr-based message; getFullOutput() exposes
140+
// both streams so the caller can check for "unknown flag: --profile" in either.
141+
throw new CliCommandException(
142+
"cannot get access token: " + stderr.trim(), stdout + "\n" + stderr);
105143
}
106144
JsonNode jsonNode = new ObjectMapper().readTree(stdout);
107145
String tokenType = jsonNode.get(tokenTypeField).asText();
@@ -111,8 +149,33 @@ public Token getToken() {
111149
return new Token(accessToken, tokenType, expiresOn);
112150
} catch (DatabricksException e) {
113151
throw e;
114-
} catch (InterruptedException | IOException e) {
115-
throw new DatabricksException("cannot get access token: " + e.getMessage(), e);
152+
} catch (InterruptedException e) {
153+
throw new IOException("cannot get access token: " + e.getMessage(), e);
154+
}
155+
}
156+
157+
@Override
158+
public Token getToken() {
159+
try {
160+
return execCliCommand(this.cmd);
161+
} catch (IOException e) {
162+
String textToCheck =
163+
e instanceof CliCommandException
164+
? ((CliCommandException) e).getFullOutput()
165+
: e.getMessage();
166+
if (fallbackCmd != null
167+
&& textToCheck != null
168+
&& textToCheck.contains("unknown flag: --profile")) {
169+
LOG.warn(
170+
"Databricks CLI does not support --profile flag. Falling back to --host. "
171+
+ "Please upgrade your CLI to the latest version.");
172+
try {
173+
return execCliCommand(this.fallbackCmd);
174+
} catch (IOException fallbackException) {
175+
throw new DatabricksException(fallbackException.getMessage(), fallbackException);
176+
}
177+
}
178+
throw new DatabricksException(e.getMessage(), e);
116179
}
117180
}
118181
}

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

Lines changed: 22 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -21,13 +21,13 @@ public String authType() {
2121
}
2222

2323
/**
24-
* Builds the CLI command arguments for the databricks auth token command.
24+
* Builds the CLI command arguments using --host (legacy path).
2525
*
2626
* @param cliPath Path to the databricks CLI executable
2727
* @param config Configuration containing host, account ID, workspace ID, etc.
2828
* @return List of command arguments
2929
*/
30-
List<String> buildCliCommand(String cliPath, DatabricksConfig config) {
30+
static List<String> buildHostArgs(String cliPath, DatabricksConfig config) {
3131
List<String> cmd =
3232
new ArrayList<>(Arrays.asList(cliPath, "auth", "token", "--host", config.getHost()));
3333
if (config.getExperimentalIsUnifiedHost() != null && config.getExperimentalIsUnifiedHost()) {
@@ -57,8 +57,26 @@ private CliTokenSource getDatabricksCliTokenSource(DatabricksConfig config) {
5757
LOG.debug("Databricks CLI could not be found");
5858
return null;
5959
}
60-
List<String> cmd = buildCliCommand(cliPath, config);
61-
return new CliTokenSource(cmd, "token_type", "access_token", "expiry", config.getEnv());
60+
61+
List<String> cmd;
62+
List<String> fallbackCmd = null;
63+
64+
if (config.getProfile() != null) {
65+
// When profile is set, use --profile as the primary command.
66+
// The profile contains the full config (host, account_id, etc.).
67+
cmd =
68+
new ArrayList<>(
69+
Arrays.asList(cliPath, "auth", "token", "--profile", config.getProfile()));
70+
// Build a --host fallback for older CLIs that don't support --profile.
71+
if (config.getHost() != null) {
72+
fallbackCmd = buildHostArgs(cliPath, config);
73+
}
74+
} else {
75+
cmd = buildHostArgs(cliPath, config);
76+
}
77+
78+
return new CliTokenSource(
79+
cmd, "token_type", "access_token", "expiry", config.getEnv(), fallbackCmd);
6280
}
6381

6482
@Override

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

Lines changed: 45 additions & 14 deletions
Original file line numberDiff line numberDiff line change
@@ -18,19 +18,19 @@ class DatabricksCliCredentialsProviderTest {
1818
private final DatabricksCliCredentialsProvider provider = new DatabricksCliCredentialsProvider();
1919

2020
@Test
21-
void testBuildCliCommand_WorkspaceHost() {
21+
void testBuildHostArgs_WorkspaceHost() {
2222
DatabricksConfig config = new DatabricksConfig().setHost(HOST);
2323

24-
List<String> cmd = provider.buildCliCommand(CLI_PATH, config);
24+
List<String> cmd = provider.buildHostArgs(CLI_PATH, config);
2525

2626
assertEquals(Arrays.asList(CLI_PATH, "auth", "token", "--host", HOST), cmd);
2727
}
2828

2929
@Test
30-
void testBuildCliCommand_AccountHost() {
30+
void testBuildHostArgs_AccountHost() {
3131
DatabricksConfig config = new DatabricksConfig().setHost(ACCOUNT_HOST).setAccountId(ACCOUNT_ID);
3232

33-
List<String> cmd = provider.buildCliCommand(CLI_PATH, config);
33+
List<String> cmd = provider.buildHostArgs(CLI_PATH, config);
3434

3535
assertEquals(
3636
Arrays.asList(
@@ -39,15 +39,15 @@ void testBuildCliCommand_AccountHost() {
3939
}
4040

4141
@Test
42-
void testBuildCliCommand_UnifiedHost_WithAccountIdAndWorkspaceId() {
42+
void testBuildHostArgs_UnifiedHost_WithAccountIdAndWorkspaceId() {
4343
DatabricksConfig config =
4444
new DatabricksConfig()
4545
.setHost(UNIFIED_HOST)
4646
.setExperimentalIsUnifiedHost(true)
4747
.setAccountId(ACCOUNT_ID)
4848
.setWorkspaceId(WORKSPACE_ID);
4949

50-
List<String> cmd = provider.buildCliCommand(CLI_PATH, config);
50+
List<String> cmd = provider.buildHostArgs(CLI_PATH, config);
5151

5252
assertEquals(
5353
Arrays.asList(
@@ -65,14 +65,14 @@ void testBuildCliCommand_UnifiedHost_WithAccountIdAndWorkspaceId() {
6565
}
6666

6767
@Test
68-
void testBuildCliCommand_UnifiedHost_WithAccountIdOnly() {
68+
void testBuildHostArgs_UnifiedHost_WithAccountIdOnly() {
6969
DatabricksConfig config =
7070
new DatabricksConfig()
7171
.setHost(UNIFIED_HOST)
7272
.setExperimentalIsUnifiedHost(true)
7373
.setAccountId(ACCOUNT_ID);
7474

75-
List<String> cmd = provider.buildCliCommand(CLI_PATH, config);
75+
List<String> cmd = provider.buildHostArgs(CLI_PATH, config);
7676

7777
assertEquals(
7878
Arrays.asList(
@@ -88,14 +88,14 @@ void testBuildCliCommand_UnifiedHost_WithAccountIdOnly() {
8888
}
8989

9090
@Test
91-
void testBuildCliCommand_UnifiedHost_WithWorkspaceIdOnly() {
91+
void testBuildHostArgs_UnifiedHost_WithWorkspaceIdOnly() {
9292
DatabricksConfig config =
9393
new DatabricksConfig()
9494
.setHost(UNIFIED_HOST)
9595
.setExperimentalIsUnifiedHost(true)
9696
.setWorkspaceId(WORKSPACE_ID);
9797

98-
List<String> cmd = provider.buildCliCommand(CLI_PATH, config);
98+
List<String> cmd = provider.buildHostArgs(CLI_PATH, config);
9999

100100
assertEquals(
101101
Arrays.asList(
@@ -111,11 +111,11 @@ void testBuildCliCommand_UnifiedHost_WithWorkspaceIdOnly() {
111111
}
112112

113113
@Test
114-
void testBuildCliCommand_UnifiedHost_WithNoAccountIdOrWorkspaceId() {
114+
void testBuildHostArgs_UnifiedHost_WithNoAccountIdOrWorkspaceId() {
115115
DatabricksConfig config =
116116
new DatabricksConfig().setHost(UNIFIED_HOST).setExperimentalIsUnifiedHost(true);
117117

118-
List<String> cmd = provider.buildCliCommand(CLI_PATH, config);
118+
List<String> cmd = provider.buildHostArgs(CLI_PATH, config);
119119

120120
assertEquals(
121121
Arrays.asList(
@@ -124,19 +124,50 @@ void testBuildCliCommand_UnifiedHost_WithNoAccountIdOrWorkspaceId() {
124124
}
125125

126126
@Test
127-
void testBuildCliCommand_UnifiedHostFalse_WithAccountHost() {
127+
void testBuildHostArgs_UnifiedHostFalse_WithAccountHost() {
128128
// When experimentalIsUnifiedHost is explicitly false, should fall back to account-id logic
129129
DatabricksConfig config =
130130
new DatabricksConfig()
131131
.setHost(ACCOUNT_HOST)
132132
.setExperimentalIsUnifiedHost(false)
133133
.setAccountId(ACCOUNT_ID);
134134

135-
List<String> cmd = provider.buildCliCommand(CLI_PATH, config);
135+
List<String> cmd = provider.buildHostArgs(CLI_PATH, config);
136136

137137
assertEquals(
138138
Arrays.asList(
139139
CLI_PATH, "auth", "token", "--host", ACCOUNT_HOST, "--account-id", ACCOUNT_ID),
140140
cmd);
141141
}
142+
143+
@Test
144+
void testProfile_UsesPrimaryProfileCmdWithHostFallback() {
145+
// When profile is set and host is present, --profile is primary and --host is fallback.
146+
// We verify this by inspecting buildHostArgs (fallback path) and the profile args (primary).
147+
DatabricksConfig config = new DatabricksConfig().setHost(HOST).setProfile("my-profile");
148+
149+
// The primary command that getDatabricksCliTokenSource would build
150+
List<String> expectedPrimary =
151+
Arrays.asList(CLI_PATH, "auth", "token", "--profile", "my-profile");
152+
// The fallback command is the host-based one
153+
List<String> expectedFallback = provider.buildHostArgs(CLI_PATH, config);
154+
155+
assertEquals(Arrays.asList(CLI_PATH, "auth", "token", "--host", HOST), expectedFallback);
156+
assertEquals(
157+
Arrays.asList("auth", "token", "--profile", "my-profile"),
158+
expectedPrimary.subList(1, expectedPrimary.size()));
159+
assertFalse(expectedFallback.contains("--profile"));
160+
assertTrue(expectedFallback.contains("--host"));
161+
}
162+
163+
@Test
164+
void testProfile_NoHostFallbackWhenHostAbsent() {
165+
// When profile is set but host is null, buildHostArgs is not called so no fallback is built.
166+
// This test confirms buildHostArgs correctly uses host when present.
167+
DatabricksConfig config = new DatabricksConfig().setHost(HOST);
168+
List<String> hostArgs = provider.buildHostArgs(CLI_PATH, config);
169+
assertTrue(hostArgs.contains("--host"));
170+
assertTrue(hostArgs.contains(HOST));
171+
assertFalse(hostArgs.contains("--profile"));
172+
}
142173
}

0 commit comments

Comments
 (0)