diff --git a/cmd/auth/token.go b/cmd/auth/token.go index 5c8e5e8771..ca8582bd02 100644 --- a/cmd/auth/token.go +++ b/cmd/auth/token.go @@ -124,6 +124,13 @@ func loadToken(ctx context.Context, args loadTokenArgs) (*oauth2.Token, error) { return nil, errors.New("providing both a profile and host is not supported") } + // When no explicit --profile flag is provided, check the env var. This + // handles the case where downstream tools (like the Terraform provider) + // pass --host but not --profile, while DATABRICKS_CONFIG_PROFILE is set. + if args.profileName == "" { + args.profileName = env.Get(ctx, "DATABRICKS_CONFIG_PROFILE") + } + // If no --profile flag, try resolving the positional arg as a profile name. // If it matches, use it. If not, fall through to host treatment. if args.profileName == "" && len(args.args) == 1 { diff --git a/cmd/auth/token_test.go b/cmd/auth/token_test.go index 8becf5f43c..191b9ca609 100644 --- a/cmd/auth/token_test.go +++ b/cmd/auth/token_test.go @@ -172,6 +172,10 @@ func TestToken_loadToken(t *testing.T) { RefreshToken: "legacy-ws", Expiry: time.Now().Add(1 * time.Hour), }, + "dup1": { + RefreshToken: "dup1", + Expiry: time.Now().Add(1 * time.Hour), + }, }, } validateToken := func(resp *oauth2.Token) { @@ -613,7 +617,7 @@ func TestToken_loadToken(t *testing.T) { validateToken: validateToken, }, { - name: "no args, DATABRICKS_HOST takes precedence over DATABRICKS_CONFIG_PROFILE", + name: "no args, DATABRICKS_CONFIG_PROFILE env takes precedence over DATABRICKS_HOST", setupCtx: func(ctx context.Context) context.Context { ctx = env.Set(ctx, "DATABRICKS_HOST", "https://workspace-a.cloud.databricks.com") ctx = env.Set(ctx, "DATABRICKS_CONFIG_PROFILE", "expired") @@ -633,6 +637,27 @@ func TestToken_loadToken(t *testing.T) { }, validateToken: validateToken, }, + { + name: "host flag with profile env var disambiguates multi-profile", + setupCtx: func(ctx context.Context) context.Context { + return env.Set(ctx, "DATABRICKS_CONFIG_PROFILE", "dup1") + }, + args: loadTokenArgs{ + authArguments: &auth.AuthArguments{ + Host: "https://shared.cloud.databricks.com", + }, + profileName: "", + args: []string{}, + tokenTimeout: 1 * time.Hour, + profiler: profiler, + persistentAuthOpts: []u2m.PersistentAuthOption{ + u2m.WithTokenCache(tokenCache), + u2m.WithOAuthEndpointSupplier(&MockApiClient{}), + u2m.WithHttpClient(&http.Client{Transport: fixtures.SliceTransport{refreshSuccessTokenResponse}}), + }, + }, + validateToken: validateToken, + }, } for _, c := range cases { t.Run(c.name, func(t *testing.T) { diff --git a/libs/auth/env.go b/libs/auth/env.go index 50ec4d0a17..ebd399e559 100644 --- a/libs/auth/env.go +++ b/libs/auth/env.go @@ -14,12 +14,6 @@ import ( func Env(cfg *config.Config) map[string]string { out := make(map[string]string) for _, attr := range config.ConfigAttributes { - // Ignore profile so that downstream tools don't try and reload - // the profile. We know the current configuration is already valid since - // otherwise the CLI would have thrown an error when loading it. - if attr.Name == "profile" { - continue - } if len(attr.EnvVars) == 0 { continue } diff --git a/libs/auth/env_test.go b/libs/auth/env_test.go index 5aa893e110..b1a41187f7 100644 --- a/libs/auth/env_test.go +++ b/libs/auth/env_test.go @@ -9,7 +9,7 @@ import ( func TestAuthEnv(t *testing.T) { in := &config.Config{ - Profile: "thisshouldbeignored", + Profile: "myprofile", Host: "https://test.com", Token: "test-token", Password: "test-password", @@ -24,6 +24,7 @@ func TestAuthEnv(t *testing.T) { } expected := map[string]string{ + "DATABRICKS_CONFIG_PROFILE": "myprofile", "DATABRICKS_HOST": "https://test.com", "DATABRICKS_TOKEN": "test-token", "DATABRICKS_PASSWORD": "test-password",