Skip to content

Commit 4cb4c26

Browse files
Generalize CLI token source into progressive command list
Replace the three explicit command fields (forceCmd, profileCmd, hostCmd) with a []cliCommand list and an activeCommandIndex. Feature flags are defined statically in cliFeatureFlags and stripped progressively to build commands for older CLI versions. Token() now iterates from activeCommandIndex, falling back on unknown flag errors and caching the working command index for subsequent calls. Adding future flags (e.g. --scopes) requires only a one-line addition to the cliFeatureFlags slice. Signed-off-by: Mihai Mitrea <mihai.mitrea@databricks.com>
1 parent f626fed commit 4cb4c26

3 files changed

Lines changed: 178 additions & 104 deletions

File tree

NEXT_CHANGELOG.md

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

2222
### Internal Changes
2323

24+
* Generalize CLI token source into a progressive command list for forward-compatible flag support.
2425
* Normalize internal token sources on `auth.TokenSource` for proper context propagation ([#1577](https://github.com/databricks/databricks-sdk-go/pull/1577)).
2526
* Bump golang.org/x/crypto from 0.21.0 to 0.45.0 in /examples/slog ([#1566](https://github.com/databricks/databricks-sdk-go/pull/1566)).
2627
* Bump golang.org/x/net from 0.23.0 to 0.33.0 in /examples/slog ([#1127](https://github.com/databricks/databricks-sdk-go/pull/1127)).

config/cli_token_source.go

Lines changed: 64 additions & 49 deletions
Original file line numberDiff line numberDiff line change
@@ -31,50 +31,78 @@ type cliTokenResponse struct {
3131
Expiry string `json:"expiry"`
3232
}
3333

34-
// CliTokenSource fetches OAuth tokens by shelling out to the Databricks CLI.
35-
// Commands are tried in order: forceCmd -> profileCmd -> hostCmd, progressively
36-
// falling back to simpler invocations for older CLI versions.
37-
type CliTokenSource struct {
38-
// forceCmd uses --profile with --force-refresh to bypass the CLI's token cache.
39-
// Nil when cfg.Profile is empty (--force-refresh requires --profile support).
40-
forceCmd []string
34+
// cliFeatureFlag defines a CLI feature flag and the warning to log when
35+
// falling back because the CLI does not support it. Ordered newest-first:
36+
// commands are built by progressively stripping these flags for older CLIs.
37+
type cliFeatureFlag struct {
38+
flag string
39+
warningMessage string
40+
}
41+
42+
var cliFeatureFlags = []cliFeatureFlag{
43+
{"--force-refresh", "Databricks CLI does not support --force-refresh flag. The CLI's token cache may provide stale tokens. Please upgrade your CLI to the latest version."},
44+
}
4145

42-
// profileCmd uses --profile for token lookup. Nil when cfg.Profile is empty.
43-
profileCmd []string
46+
const profileFlagWarning = "Databricks CLI does not support --profile flag. Falling back to --host. Please upgrade your CLI to the latest version."
4447

45-
// hostCmd uses --host as a fallback for CLIs that predate --profile support.
46-
// Nil when cfg.Host is empty.
47-
hostCmd []string
48+
// cliCommand is a single CLI invocation with an associated warning message
49+
// that is logged when this command fails and we fall back to the next one.
50+
type cliCommand struct {
51+
args []string
52+
warningMessage string
53+
}
54+
55+
// CliTokenSource fetches OAuth tokens by shelling out to the Databricks CLI.
56+
// It holds a list of commands ordered from most feature-rich to simplest,
57+
// falling back progressively for older CLI versions that lack newer flags.
58+
type CliTokenSource struct {
59+
commands []cliCommand
60+
activeCommandIndex int
4861
}
4962

5063
func NewCliTokenSource(cfg *Config) (*CliTokenSource, error) {
5164
cliPath, err := findDatabricksCli(cfg.DatabricksCliPath)
5265
if err != nil {
5366
return nil, err
5467
}
55-
forceCmd, profileCmd, hostCmd := buildCliCommands(cliPath, cfg)
56-
return &CliTokenSource{forceCmd: forceCmd, profileCmd: profileCmd, hostCmd: hostCmd}, nil
68+
return &CliTokenSource{commands: buildCliCommands(cliPath, cfg)}, nil
5769
}
5870

59-
// buildCliCommands constructs the CLI commands for fetching an auth token.
60-
// When cfg.Profile is set, three commands are built: a --force-refresh variant,
61-
// a plain --profile variant, and (when host is available) a --host fallback.
62-
// When cfg.Profile is empty, only --host is returned — the CLI must support
63-
// --profile before --force-refresh can be used (monotonic feature assumption).
64-
func buildCliCommands(cliPath string, cfg *Config) ([]string, []string, []string) {
65-
var forceCmd, profileCmd, hostCmd []string
71+
// buildCliCommands constructs the list of CLI commands for fetching an auth
72+
// token, ordered from most feature-rich to simplest. When cfg.Profile is set,
73+
// commands include feature flags from [cliFeatureFlags] (stripped progressively)
74+
// plus a --host fallback when host is available. When cfg.Profile is empty,
75+
// only --host is returned — the CLI must support --profile before any feature
76+
// flags can be used (monotonic feature assumption).
77+
func buildCliCommands(cliPath string, cfg *Config) []cliCommand {
78+
var commands []cliCommand
6679
if cfg.Profile != "" {
67-
profileCmd = []string{cliPath, "auth", "token", "--profile", cfg.Profile}
68-
forceCmd = append(profileCmd, "--force-refresh")
80+
baseArgs := []string{cliPath, "auth", "token", "--profile", cfg.Profile}
81+
for i := 0; i <= len(cliFeatureFlags); i++ {
82+
args := append([]string{}, baseArgs...)
83+
for _, f := range cliFeatureFlags[i:] {
84+
args = append(args, f.flag)
85+
}
86+
warning := profileFlagWarning
87+
if i < len(cliFeatureFlags) {
88+
warning = cliFeatureFlags[i].warningMessage
89+
}
90+
commands = append(commands, cliCommand{args: args, warningMessage: warning})
91+
}
6992
}
70-
if cfg.Host != "" {
71-
hostCmd = buildHostCommand(cliPath, cfg)
93+
hostArgs := buildHostCommand(cliPath, cfg)
94+
if hostArgs != nil {
95+
commands = append(commands, cliCommand{args: hostArgs})
7296
}
73-
return forceCmd, profileCmd, hostCmd
97+
return commands
7498
}
7599

76100
// buildHostCommand constructs the legacy --host based CLI command.
101+
// Returns nil when cfg.Host is empty.
77102
func buildHostCommand(cliPath string, cfg *Config) []string {
103+
if cfg.Host == "" {
104+
return nil
105+
}
78106
cmd := []string{cliPath, "auth", "token", "--host", cfg.Host}
79107
switch cfg.HostType() {
80108
case AccountHost:
@@ -84,36 +112,23 @@ func buildHostCommand(cliPath string, cfg *Config) []string {
84112
}
85113

86114
// Token fetches an OAuth token by shelling out to the Databricks CLI.
87-
// Commands are tried in order: forceCmd -> profileCmd -> hostCmd, skipping nil
88-
// entries. Each command falls through to the next on "unknown flag" errors,
89-
// logging a warning about the unsupported feature.
115+
// Commands are tried from [activeCommandIndex] forward, falling back to simpler
116+
// commands when the CLI reports an unknown flag. Once a command succeeds,
117+
// [activeCommandIndex] is updated so subsequent calls skip commands known to fail.
90118
func (c *CliTokenSource) Token(ctx context.Context) (*oauth2.Token, error) {
91-
if c.forceCmd != nil {
92-
tok, err := c.execCliCommand(ctx, c.forceCmd)
119+
for i := c.activeCommandIndex; i < len(c.commands); i++ {
120+
cmd := c.commands[i]
121+
tok, err := c.execCliCommand(ctx, cmd.args)
93122
if err == nil {
123+
c.activeCommandIndex = i
94124
return tok, nil
95125
}
96-
if !isUnknownFlagError(err, "") {
126+
lastCommand := i == len(c.commands)-1
127+
if lastCommand || !isUnknownFlagError(err, "") {
97128
return nil, err
98129
}
99-
logger.Warnf(ctx, "Databricks CLI does not support --force-refresh flag. The CLI's token cache may provide stale tokens. Please upgrade your CLI to the latest version.")
130+
logger.Warnf(ctx, cmd.warningMessage)
100131
}
101-
102-
if c.profileCmd != nil {
103-
tok, err := c.execCliCommand(ctx, c.profileCmd)
104-
if err == nil {
105-
return tok, nil
106-
}
107-
if !isUnknownFlagError(err, "--profile") {
108-
return nil, err
109-
}
110-
logger.Warnf(ctx, "Databricks CLI does not support --profile flag. Falling back to --host. Please upgrade your CLI to the latest version.")
111-
}
112-
113-
if c.hostCmd != nil {
114-
return c.execCliCommand(ctx, c.hostCmd)
115-
}
116-
117132
return nil, fmt.Errorf("no CLI command configured")
118133
}
119134

0 commit comments

Comments
 (0)