Skip to content

Commit 6b8a57f

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 6b8a57f

File tree

5 files changed

+294
-28
lines changed

5 files changed

+294
-28
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: 57 additions & 17 deletions
Original file line numberDiff line numberDiff line change
@@ -25,15 +25,17 @@
2525
public class CliTokenSource implements TokenSource {
2626
private static final Logger LOG = LoggerFactory.getLogger(CliTokenSource.class);
2727

28-
private List<String> cmd;
28+
// forceCmd is tried before profileCmd when non-null. If the CLI rejects
29+
// --force-refresh or --profile, execution falls through to profileCmd.
30+
private List<String> forceCmd;
31+
32+
private List<String> profileCmd;
2933
private String tokenTypeField;
3034
private String accessTokenField;
3135
private String expiryField;
3236
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+
// fallbackCmd is tried when profileCmd fails with "unknown flag: --profile",
38+
// indicating the CLI is too old to support --profile.
3739
private List<String> fallbackCmd;
3840

3941
/**
@@ -58,7 +60,7 @@ public CliTokenSource(
5860
String accessTokenField,
5961
String expiryField,
6062
Environment env) {
61-
this(cmd, tokenTypeField, accessTokenField, expiryField, env, null);
63+
this(cmd, tokenTypeField, accessTokenField, expiryField, env, null, null);
6264
}
6365

6466
public CliTokenSource(
@@ -68,14 +70,26 @@ public CliTokenSource(
6870
String expiryField,
6971
Environment env,
7072
List<String> fallbackCmd) {
73+
this(cmd, tokenTypeField, accessTokenField, expiryField, env, fallbackCmd, null);
74+
}
75+
76+
public CliTokenSource(
77+
List<String> cmd,
78+
String tokenTypeField,
79+
String accessTokenField,
80+
String expiryField,
81+
Environment env,
82+
List<String> fallbackCmd,
83+
List<String> forceCmd) {
7184
super();
72-
this.cmd = OSUtils.get(env).getCliExecutableCommand(cmd);
85+
this.profileCmd = OSUtils.get(env).getCliExecutableCommand(cmd);
7386
this.tokenTypeField = tokenTypeField;
7487
this.accessTokenField = accessTokenField;
7588
this.expiryField = expiryField;
7689
this.env = env;
7790
this.fallbackCmd =
7891
fallbackCmd != null ? OSUtils.get(env).getCliExecutableCommand(fallbackCmd) : null;
92+
this.forceCmd = forceCmd != null ? OSUtils.get(env).getCliExecutableCommand(forceCmd) : null;
7993
}
8094

8195
/**
@@ -153,18 +167,22 @@ private Token execCliCommand(List<String> cmdToRun) throws IOException {
153167
}
154168
}
155169

156-
@Override
157-
public Token getToken() {
170+
private String getErrorText(IOException e) {
171+
return e instanceof CliCommandException
172+
? ((CliCommandException) e).getFullOutput()
173+
: e.getMessage();
174+
}
175+
176+
private boolean isUnknownFlagError(String errorText, String flag) {
177+
return errorText != null && errorText.contains("unknown flag: " + flag);
178+
}
179+
180+
private Token execProfileCmdWithFallback() {
158181
try {
159-
return execCliCommand(this.cmd);
182+
return execCliCommand(this.profileCmd);
160183
} 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")) {
184+
String textToCheck = getErrorText(e);
185+
if (fallbackCmd != null && isUnknownFlagError(textToCheck, "--profile")) {
168186
LOG.warn(
169187
"Databricks CLI does not support --profile flag. Falling back to --host. "
170188
+ "Please upgrade your CLI to the latest version.");
@@ -177,4 +195,26 @@ public Token getToken() {
177195
throw new DatabricksException(e.getMessage(), e);
178196
}
179197
}
198+
199+
@Override
200+
public Token getToken() {
201+
if (forceCmd == null) {
202+
return execProfileCmdWithFallback();
203+
}
204+
205+
try {
206+
return execCliCommand(this.forceCmd);
207+
} catch (IOException e) {
208+
String textToCheck = getErrorText(e);
209+
if (isUnknownFlagError(textToCheck, "--force-refresh")
210+
|| isUnknownFlagError(textToCheck, "--profile")) {
211+
LOG.warn(
212+
"Databricks CLI does not support --force-refresh flag. "
213+
+ "Falling back to regular token fetch. "
214+
+ "Please upgrade your CLI to the latest version.");
215+
return execProfileCmdWithFallback();
216+
}
217+
throw new DatabricksException(e.getMessage(), e);
218+
}
219+
}
180220
}

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

Lines changed: 18 additions & 9 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) {
@@ -79,25 +90,23 @@ private CliTokenSource getDatabricksCliTokenSource(DatabricksConfig config) {
7990
return null;
8091
}
8192

82-
List<String> cmd;
93+
List<String> profileCmd;
8394
List<String> fallbackCmd = null;
95+
List<String> forceCmd;
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+
profileCmd = buildProfileArgs(cliPath, config);
99+
forceCmd = withForceRefresh(profileCmd);
92100
if (config.getHost() != null) {
93101
fallbackCmd = buildHostArgs(cliPath, config);
94102
}
95103
} else {
96-
cmd = buildHostArgs(cliPath, config);
104+
profileCmd = buildHostArgs(cliPath, config);
105+
forceCmd = withForceRefresh(profileCmd);
97106
}
98107

99108
return new CliTokenSource(
100-
cmd, "token_type", "access_token", "expiry", config.getEnv(), fallbackCmd);
109+
profileCmd, "token_type", "access_token", "expiry", config.getEnv(), fallbackCmd, forceCmd);
101110
}
102111

103112
@Override

0 commit comments

Comments
 (0)