Skip to content

Commit dd71f92

Browse files
hectorcast-dbclaude
hectorcast-db
and
claude
committed
Pass --profile to CLI token source and add --host fallback for older CLIs
Port of databricks-sdk-py#1297 / databricks-sdk-go#1497. - CliTokenSource: add optional fallbackCmd; extract execCliCommand(); on IOError check full stdout+stderr for "unknown flag: --profile" via CliCommandException (clean stderr message exposed to users, full output used only for detection); if matched, warn and retry with fallbackCmd. - DatabricksCliCredentialsProvider: rename buildCliCommand -> buildHostArgs; when cfg.profile is set use --profile as primary command, build --host fallback if cfg.host is also present; when profile is absent keep existing --host path. Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
1 parent d98ea37 commit dd71f92

File tree

3 files changed

+178
-32
lines changed

3 files changed

+178
-32
lines changed

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

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -27,7 +27,7 @@ public String authType() {
2727
* @param config Configuration containing host, account ID, workspace ID, etc.
2828
* @return List of command arguments
2929
*/
30-
static List<String> buildHostArgs(String cliPath, DatabricksConfig config) {
30+
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()) {

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

Lines changed: 177 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -2,6 +2,7 @@
22

33
import static org.junit.jupiter.api.Assertions.assertEquals;
44
import static org.junit.jupiter.api.Assertions.assertThrows;
5+
import static org.junit.jupiter.api.Assertions.assertTrue;
56
import static org.mockito.ArgumentMatchers.any;
67
import static org.mockito.Mockito.mock;
78
import static org.mockito.Mockito.mockConstruction;
@@ -27,9 +28,11 @@
2728
import java.util.List;
2829
import java.util.Map;
2930
import java.util.TimeZone;
31+
import java.util.concurrent.atomic.AtomicInteger;
3032
import java.util.stream.Collectors;
3133
import java.util.stream.IntStream;
3234
import java.util.stream.Stream;
35+
import org.junit.jupiter.api.Test;
3336
import org.junit.jupiter.params.ParameterizedTest;
3437
import org.junit.jupiter.params.provider.Arguments;
3538
import org.junit.jupiter.params.provider.MethodSource;
@@ -213,4 +216,178 @@ public void testParseExpiry(String input, Instant expectedInstant, String descri
213216
assertEquals(expectedInstant, parsedInstant);
214217
}
215218
}
219+
220+
// ---- Fallback tests for --profile flag handling ----
221+
222+
private CliTokenSource makeTokenSource(
223+
Environment env, List<String> primaryCmd, List<String> fallbackCmd) {
224+
OSUtilities osUtils = mock(OSUtilities.class);
225+
when(osUtils.getCliExecutableCommand(any())).thenAnswer(inv -> inv.getArgument(0));
226+
try (MockedStatic<OSUtils> mockedOSUtils = mockStatic(OSUtils.class)) {
227+
mockedOSUtils.when(() -> OSUtils.get(any())).thenReturn(osUtils);
228+
return new CliTokenSource(
229+
primaryCmd, "token_type", "access_token", "expiry", env, fallbackCmd);
230+
}
231+
}
232+
233+
private String validTokenJson(String accessToken) {
234+
String expiry =
235+
ZonedDateTime.now()
236+
.plusHours(1)
237+
.format(DateTimeFormatter.ofPattern("yyyy-MM-dd'T'HH:mm:ss.SSSXXX"));
238+
return String.format(
239+
"{\"token_type\":\"Bearer\",\"access_token\":\"%s\",\"expiry\":\"%s\"}",
240+
accessToken, expiry);
241+
}
242+
243+
@Test
244+
public void testFallbackOnUnknownProfileFlagInStderr() {
245+
Environment env = mock(Environment.class);
246+
when(env.getEnv()).thenReturn(new HashMap<>());
247+
248+
List<String> primaryCmd =
249+
Arrays.asList("databricks", "auth", "token", "--profile", "my-profile");
250+
List<String> fallbackCmdList =
251+
Arrays.asList("databricks", "auth", "token", "--host", "https://workspace.databricks.com");
252+
253+
CliTokenSource tokenSource = makeTokenSource(env, primaryCmd, fallbackCmdList);
254+
255+
AtomicInteger callCount = new AtomicInteger(0);
256+
try (MockedConstruction<ProcessBuilder> mocked =
257+
mockConstruction(
258+
ProcessBuilder.class,
259+
(pb, context) -> {
260+
if (callCount.getAndIncrement() == 0) {
261+
Process failProcess = mock(Process.class);
262+
when(failProcess.getInputStream())
263+
.thenReturn(new ByteArrayInputStream(new byte[0]));
264+
when(failProcess.getErrorStream())
265+
.thenReturn(
266+
new ByteArrayInputStream("Error: unknown flag: --profile".getBytes()));
267+
when(failProcess.waitFor()).thenReturn(1);
268+
when(pb.start()).thenReturn(failProcess);
269+
} else {
270+
Process successProcess = mock(Process.class);
271+
when(successProcess.getInputStream())
272+
.thenReturn(new ByteArrayInputStream(validTokenJson("fallback-token").getBytes()));
273+
when(successProcess.getErrorStream())
274+
.thenReturn(new ByteArrayInputStream(new byte[0]));
275+
when(successProcess.waitFor()).thenReturn(0);
276+
when(pb.start()).thenReturn(successProcess);
277+
}
278+
})) {
279+
Token token = tokenSource.getToken();
280+
assertEquals("fallback-token", token.getAccessToken());
281+
assertEquals(2, mocked.constructed().size());
282+
}
283+
}
284+
285+
@Test
286+
public void testFallbackTriggeredWhenUnknownFlagInStdout() {
287+
// Fallback triggers even when "unknown flag" appears in stdout rather than stderr.
288+
Environment env = mock(Environment.class);
289+
when(env.getEnv()).thenReturn(new HashMap<>());
290+
291+
List<String> primaryCmd =
292+
Arrays.asList("databricks", "auth", "token", "--profile", "my-profile");
293+
List<String> fallbackCmdList =
294+
Arrays.asList("databricks", "auth", "token", "--host", "https://workspace.databricks.com");
295+
296+
CliTokenSource tokenSource = makeTokenSource(env, primaryCmd, fallbackCmdList);
297+
298+
AtomicInteger callCount = new AtomicInteger(0);
299+
try (MockedConstruction<ProcessBuilder> mocked =
300+
mockConstruction(
301+
ProcessBuilder.class,
302+
(pb, context) -> {
303+
if (callCount.getAndIncrement() == 0) {
304+
Process failProcess = mock(Process.class);
305+
when(failProcess.getInputStream())
306+
.thenReturn(
307+
new ByteArrayInputStream(
308+
"Error: unknown flag: --profile".getBytes()));
309+
when(failProcess.getErrorStream())
310+
.thenReturn(new ByteArrayInputStream(new byte[0]));
311+
when(failProcess.waitFor()).thenReturn(1);
312+
when(pb.start()).thenReturn(failProcess);
313+
} else {
314+
Process successProcess = mock(Process.class);
315+
when(successProcess.getInputStream())
316+
.thenReturn(new ByteArrayInputStream(validTokenJson("fallback-token").getBytes()));
317+
when(successProcess.getErrorStream())
318+
.thenReturn(new ByteArrayInputStream(new byte[0]));
319+
when(successProcess.waitFor()).thenReturn(0);
320+
when(pb.start()).thenReturn(successProcess);
321+
}
322+
})) {
323+
Token token = tokenSource.getToken();
324+
assertEquals("fallback-token", token.getAccessToken());
325+
assertEquals(2, mocked.constructed().size());
326+
}
327+
}
328+
329+
@Test
330+
public void testNoFallbackOnRealAuthError() {
331+
// When the primary fails with a real error (not unknown flag), no fallback is attempted.
332+
Environment env = mock(Environment.class);
333+
when(env.getEnv()).thenReturn(new HashMap<>());
334+
335+
List<String> primaryCmd =
336+
Arrays.asList("databricks", "auth", "token", "--profile", "my-profile");
337+
List<String> fallbackCmdList =
338+
Arrays.asList("databricks", "auth", "token", "--host", "https://workspace.databricks.com");
339+
340+
CliTokenSource tokenSource = makeTokenSource(env, primaryCmd, fallbackCmdList);
341+
342+
try (MockedConstruction<ProcessBuilder> mocked =
343+
mockConstruction(
344+
ProcessBuilder.class,
345+
(pb, context) -> {
346+
Process failProcess = mock(Process.class);
347+
when(failProcess.getInputStream())
348+
.thenReturn(new ByteArrayInputStream(new byte[0]));
349+
when(failProcess.getErrorStream())
350+
.thenReturn(
351+
new ByteArrayInputStream(
352+
"databricks OAuth is not configured for this host".getBytes()));
353+
when(failProcess.waitFor()).thenReturn(1);
354+
when(pb.start()).thenReturn(failProcess);
355+
})) {
356+
DatabricksException ex =
357+
assertThrows(DatabricksException.class, tokenSource::getToken);
358+
assertTrue(ex.getMessage().contains("databricks OAuth is not configured"));
359+
assertEquals(1, mocked.constructed().size());
360+
}
361+
}
362+
363+
@Test
364+
public void testNoFallbackWhenFallbackCmdNotSet() {
365+
// When fallbackCmd is null and the primary fails with unknown flag, original error propagates.
366+
Environment env = mock(Environment.class);
367+
when(env.getEnv()).thenReturn(new HashMap<>());
368+
369+
List<String> primaryCmd =
370+
Arrays.asList("databricks", "auth", "token", "--profile", "my-profile");
371+
372+
CliTokenSource tokenSource = makeTokenSource(env, primaryCmd, null);
373+
374+
try (MockedConstruction<ProcessBuilder> mocked =
375+
mockConstruction(
376+
ProcessBuilder.class,
377+
(pb, context) -> {
378+
Process failProcess = mock(Process.class);
379+
when(failProcess.getInputStream())
380+
.thenReturn(new ByteArrayInputStream(new byte[0]));
381+
when(failProcess.getErrorStream())
382+
.thenReturn(
383+
new ByteArrayInputStream("Error: unknown flag: --profile".getBytes()));
384+
when(failProcess.waitFor()).thenReturn(1);
385+
when(pb.start()).thenReturn(failProcess);
386+
})) {
387+
DatabricksException ex =
388+
assertThrows(DatabricksException.class, tokenSource::getToken);
389+
assertTrue(ex.getMessage().contains("unknown flag: --profile"));
390+
assertEquals(1, mocked.constructed().size());
391+
}
392+
}
216393
}

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

Lines changed: 0 additions & 31 deletions
Original file line numberDiff line numberDiff line change
@@ -139,35 +139,4 @@ void testBuildHostArgs_UnifiedHostFalse_WithAccountHost() {
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-
}
173142
}

0 commit comments

Comments
 (0)