Skip to content

Commit 61686da

Browse files
Add --force-refresh support for Databricks CLI token fetching
Try `--force-refresh` before the regular CLI command so the SDK can bypass the CLI's own token cache when the SDK considers its token stale. If the CLI is too old to recognise `--force-refresh` (or `--profile`), gracefully fall back to the next command in the chain. Chain order: - with profile: forceCmd (--profile --force-refresh) -> profileCmd (--profile) -> fallbackCmd (--host) - without profile: forceCmd (--host --force-refresh) -> profileCmd (--host) Azure CLI callers are unchanged; they use constructors that leave forceCmd null, preserving existing behavior. Signed-off-by: Mihai Mitrea <mihai.mitrea@databricks.com>
1 parent f28430b commit 61686da

File tree

5 files changed

+285
-32
lines changed

5 files changed

+285
-32
lines changed

NEXT_CHANGELOG.md

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -4,6 +4,7 @@
44

55
### New Features and Improvements
66
* Added automatic detection of AI coding agents (Antigravity, Claude Code, Cline, Codex, Copilot CLI, Cursor, Gemini CLI, OpenCode) in the user-agent string. The SDK now appends `agent/<name>` to HTTP request headers when running inside a known AI agent environment.
7+
* Pass `--force-refresh` to Databricks CLI `auth token` command so the SDK always receives a fresh token instead of a potentially stale one from the CLI's internal cache. Falls back gracefully on older CLIs that do not support this flag.
78

89
### Bug Fixes
910
* Fixed Databricks CLI authentication to detect when the cached token's scopes don't match the SDK's configured scopes. Previously, a scope mismatch was silently ignored, causing requests to use wrong permissions. The SDK now raises an error with instructions to re-authenticate.

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

Lines changed: 42 additions & 21 deletions
Original file line numberDiff line numberDiff line change
@@ -26,15 +26,12 @@ public class CliTokenSource implements TokenSource {
2626
private static final Logger LOG = LoggerFactory.getLogger(CliTokenSource.class);
2727

2828
private List<String> cmd;
29+
private List<String> fallbackCmd;
30+
private List<String> secondFallbackCmd;
2931
private String tokenTypeField;
3032
private String accessTokenField;
3133
private String expiryField;
3234
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;
3835

3936
/**
4037
* Internal exception that carries the clean stderr message but exposes full output for checks.
@@ -58,7 +55,7 @@ public CliTokenSource(
5855
String accessTokenField,
5956
String expiryField,
6057
Environment env) {
61-
this(cmd, tokenTypeField, accessTokenField, expiryField, env, null);
58+
this(cmd, tokenTypeField, accessTokenField, expiryField, env, null, null);
6259
}
6360

6461
public CliTokenSource(
@@ -67,15 +64,19 @@ public CliTokenSource(
6764
String accessTokenField,
6865
String expiryField,
6966
Environment env,
70-
List<String> fallbackCmd) {
71-
super();
67+
List<String> fallbackCmd,
68+
List<String> secondFallbackCmd) {
7269
this.cmd = OSUtils.get(env).getCliExecutableCommand(cmd);
7370
this.tokenTypeField = tokenTypeField;
7471
this.accessTokenField = accessTokenField;
7572
this.expiryField = expiryField;
7673
this.env = env;
7774
this.fallbackCmd =
7875
fallbackCmd != null ? OSUtils.get(env).getCliExecutableCommand(fallbackCmd) : null;
76+
this.secondFallbackCmd =
77+
secondFallbackCmd != null
78+
? OSUtils.get(env).getCliExecutableCommand(secondFallbackCmd)
79+
: null;
7980
}
8081

8182
/**
@@ -153,27 +154,47 @@ private Token execCliCommand(List<String> cmdToRun) throws IOException {
153154
}
154155
}
155156

157+
private String getErrorText(IOException e) {
158+
return e instanceof CliCommandException
159+
? ((CliCommandException) e).getFullOutput()
160+
: e.getMessage();
161+
}
162+
163+
private boolean isUnknownFlagError(String errorText) {
164+
return errorText != null && errorText.contains("unknown flag:");
165+
}
166+
156167
@Override
157168
public Token getToken() {
158169
try {
159170
return execCliCommand(this.cmd);
160171
} catch (IOException e) {
161-
String textToCheck =
162-
e instanceof CliCommandException
163-
? ((CliCommandException) e).getFullOutput()
164-
: e.getMessage();
165-
if (fallbackCmd != null
166-
&& textToCheck != null
167-
&& textToCheck.contains("unknown flag: --profile")) {
172+
if (fallbackCmd != null && isUnknownFlagError(getErrorText(e))) {
168173
LOG.warn(
169-
"Databricks CLI does not support --profile flag. Falling back to --host. "
174+
"CLI does not support some flags used by this SDK. "
175+
+ "Falling back to a compatible command. "
170176
+ "Please upgrade your CLI to the latest version.");
171-
try {
172-
return execCliCommand(this.fallbackCmd);
173-
} catch (IOException fallbackException) {
174-
throw new DatabricksException(fallbackException.getMessage(), fallbackException);
175-
}
177+
} else {
178+
throw new DatabricksException(e.getMessage(), e);
179+
}
180+
}
181+
182+
try {
183+
return execCliCommand(this.fallbackCmd);
184+
} catch (IOException e) {
185+
if (secondFallbackCmd != null && isUnknownFlagError(getErrorText(e))) {
186+
LOG.warn(
187+
"CLI does not support some flags used by this SDK. "
188+
+ "Falling back to a compatible command. "
189+
+ "Please upgrade your CLI to the latest version.");
190+
} else {
191+
throw new DatabricksException(e.getMessage(), e);
176192
}
193+
}
194+
195+
try {
196+
return execCliCommand(this.secondFallbackCmd);
197+
} catch (IOException e) {
177198
throw new DatabricksException(e.getMessage(), e);
178199
}
179200
}

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

Lines changed: 23 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -69,6 +69,17 @@ List<String> buildHostArgs(String cliPath, DatabricksConfig config) {
6969
return cmd;
7070
}
7171

72+
List<String> buildProfileArgs(String cliPath, DatabricksConfig config) {
73+
return new ArrayList<>(
74+
Arrays.asList(cliPath, "auth", "token", "--profile", config.getProfile()));
75+
}
76+
77+
private static List<String> withForceRefresh(List<String> cmd) {
78+
List<String> forceCmd = new ArrayList<>(cmd);
79+
forceCmd.add("--force-refresh");
80+
return forceCmd;
81+
}
82+
7283
private CliTokenSource getDatabricksCliTokenSource(DatabricksConfig config) {
7384
String cliPath = config.getDatabricksCliPath();
7485
if (cliPath == null) {
@@ -81,23 +92,27 @@ private CliTokenSource getDatabricksCliTokenSource(DatabricksConfig config) {
8192

8293
List<String> cmd;
8394
List<String> fallbackCmd = null;
95+
List<String> secondFallbackCmd = null;
8496

8597
if (config.getProfile() != null) {
86-
// When profile is set, use --profile as the primary command.
87-
// The profile contains the full config (host, account_id, etc.).
88-
cmd =
89-
new ArrayList<>(
90-
Arrays.asList(cliPath, "auth", "token", "--profile", config.getProfile()));
91-
// Build a --host fallback for older CLIs that don't support --profile.
98+
List<String> profileArgs = buildProfileArgs(cliPath, config);
99+
cmd = withForceRefresh(profileArgs);
100+
fallbackCmd = profileArgs;
92101
if (config.getHost() != null) {
93-
fallbackCmd = buildHostArgs(cliPath, config);
102+
secondFallbackCmd = buildHostArgs(cliPath, config);
94103
}
95104
} else {
96105
cmd = buildHostArgs(cliPath, config);
97106
}
98107

99108
return new CliTokenSource(
100-
cmd, "token_type", "access_token", "expiry", config.getEnv(), fallbackCmd);
109+
cmd,
110+
"token_type",
111+
"access_token",
112+
"expiry",
113+
config.getEnv(),
114+
fallbackCmd,
115+
secondFallbackCmd);
101116
}
102117

103118
@Override

0 commit comments

Comments
 (0)