Skip to content

Commit 5e8f476

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 5e8f476

File tree

5 files changed

+297
-28
lines changed

5 files changed

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

28-
private List<String> cmd;
28+
private static final String UNKNOWN_PROFILE_FLAG = "unknown flag: --profile";
29+
private static final String UNKNOWN_FORCE_REFRESH_FLAG = "unknown flag: --force-refresh";
30+
31+
// forceCmd is tried before profileCmd when non-null. If the CLI rejects
32+
// --force-refresh or --profile, execution falls through to profileCmd.
33+
private List<String> forceCmd;
34+
35+
private List<String> profileCmd;
2936
private String tokenTypeField;
3037
private String accessTokenField;
3138
private String expiryField;
3239
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
40+
// fallbackCmd is tried when profileCmd fails with "unknown flag: --profile",
41+
// indicating the CLI is too old to support --profile.
3742
private List<String> fallbackCmd;
3843

3944
/**
@@ -58,7 +63,7 @@ public CliTokenSource(
5863
String accessTokenField,
5964
String expiryField,
6065
Environment env) {
61-
this(cmd, tokenTypeField, accessTokenField, expiryField, env, null);
66+
this(cmd, tokenTypeField, accessTokenField, expiryField, env, null, null);
6267
}
6368

6469
public CliTokenSource(
@@ -68,14 +73,26 @@ public CliTokenSource(
6873
String expiryField,
6974
Environment env,
7075
List<String> fallbackCmd) {
76+
this(cmd, tokenTypeField, accessTokenField, expiryField, env, fallbackCmd, null);
77+
}
78+
79+
public CliTokenSource(
80+
List<String> cmd,
81+
String tokenTypeField,
82+
String accessTokenField,
83+
String expiryField,
84+
Environment env,
85+
List<String> fallbackCmd,
86+
List<String> forceCmd) {
7187
super();
72-
this.cmd = OSUtils.get(env).getCliExecutableCommand(cmd);
88+
this.profileCmd = OSUtils.get(env).getCliExecutableCommand(cmd);
7389
this.tokenTypeField = tokenTypeField;
7490
this.accessTokenField = accessTokenField;
7591
this.expiryField = expiryField;
7692
this.env = env;
7793
this.fallbackCmd =
7894
fallbackCmd != null ? OSUtils.get(env).getCliExecutableCommand(fallbackCmd) : null;
95+
this.forceCmd = forceCmd != null ? OSUtils.get(env).getCliExecutableCommand(forceCmd) : null;
7996
}
8097

8198
/**
@@ -153,18 +170,22 @@ private Token execCliCommand(List<String> cmdToRun) throws IOException {
153170
}
154171
}
155172

156-
@Override
157-
public Token getToken() {
173+
private String getErrorText(IOException e) {
174+
return e instanceof CliCommandException
175+
? ((CliCommandException) e).getFullOutput()
176+
: e.getMessage();
177+
}
178+
179+
private boolean isUnknownFlagError(String errorText, String flag) {
180+
return errorText != null && errorText.contains(flag);
181+
}
182+
183+
private Token execProfileCmdWithFallback() {
158184
try {
159-
return execCliCommand(this.cmd);
185+
return execCliCommand(this.profileCmd);
160186
} 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")) {
187+
String textToCheck = getErrorText(e);
188+
if (fallbackCmd != null && isUnknownFlagError(textToCheck, UNKNOWN_PROFILE_FLAG)) {
168189
LOG.warn(
169190
"Databricks CLI does not support --profile flag. Falling back to --host. "
170191
+ "Please upgrade your CLI to the latest version.");
@@ -177,4 +198,26 @@ public Token getToken() {
177198
throw new DatabricksException(e.getMessage(), e);
178199
}
179200
}
201+
202+
@Override
203+
public Token getToken() {
204+
if (forceCmd == null) {
205+
return execProfileCmdWithFallback();
206+
}
207+
208+
try {
209+
return execCliCommand(this.forceCmd);
210+
} catch (IOException e) {
211+
String textToCheck = getErrorText(e);
212+
if (isUnknownFlagError(textToCheck, UNKNOWN_FORCE_REFRESH_FLAG)
213+
|| isUnknownFlagError(textToCheck, UNKNOWN_PROFILE_FLAG)) {
214+
LOG.warn(
215+
"Databricks CLI does not support --force-refresh flag. "
216+
+ "Falling back to regular token fetch. "
217+
+ "Please upgrade your CLI to the latest version.");
218+
return execProfileCmdWithFallback();
219+
}
220+
throw new DatabricksException(e.getMessage(), e);
221+
}
222+
}
180223
}

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)