Skip to content

Commit 9b5df04

Browse files
jh-dbjh-db
committed
address comments
1 parent f3ec955 commit 9b5df04

File tree

6 files changed

+25
-31
lines changed

6 files changed

+25
-31
lines changed

‎cmd/auth/auth.go‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -58,7 +58,7 @@ func promptForAccountID(ctx context.Context) (string, error) {
5858
}
5959

6060
func promptForWorkspaceId(ctx context.Context) (string, error) {
61-
if !cmdio.IsInTTY(ctx) {
61+
if !cmdio.IsPromptSupported(ctx) {
6262
return "", nil
6363
}
6464

‎cmd/auth/login.go‎

Lines changed: 6 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -132,7 +132,11 @@ depends on the existing profiles you have set in your configuration file
132132
if err != nil {
133133
return err
134134
}
135-
err = setHostAndAccountId(ctx, cmd, existingProfile, authArguments, args)
135+
// Set IsUnifiedHost from the profile if the flag wasn't explicitly set.
136+
if !cmd.Flag("experimental-is-unified-host").Changed && existingProfile != nil {
137+
authArguments.IsUnifiedHost = existingProfile.Experimental_IsUnifiedHost
138+
}
139+
err = setHostAndAccountId(ctx, existingProfile, authArguments, args)
136140
if err != nil {
137141
return err
138142
}
@@ -241,7 +245,7 @@ depends on the existing profiles you have set in your configuration file
241245
// 1. --account-id flag.
242246
// 2. account-id from the specified profile, if available.
243247
// 3. Prompt the user for the account-id.
244-
func setHostAndAccountId(ctx context.Context, cmd *cobra.Command, existingProfile *profile.Profile, authArguments *auth.AuthArguments, args []string) error {
248+
func setHostAndAccountId(ctx context.Context, existingProfile *profile.Profile, authArguments *auth.AuthArguments, args []string) error {
245249
// If both [HOST] and --host are provided, return an error.
246250
host := authArguments.Host
247251
if len(args) > 0 && host != "" {
@@ -267,16 +271,6 @@ func setHostAndAccountId(ctx context.Context, cmd *cobra.Command, existingProfil
267271
}
268272
}
269273

270-
// Determine if the host is a unified host in the following order of precedence:
271-
// 1. --experimental-is-unified-host flag (if explicitly set)
272-
// 2. experimental_is_unified_host from the specified profile, if available
273-
// 3. default to false if neither is provided.
274-
if !cmd.Flag("experimental-is-unified-host").Changed {
275-
if existingProfile != nil {
276-
authArguments.IsUnifiedHost = existingProfile.Experimental_IsUnifiedHost
277-
}
278-
}
279-
280274
// If the account-id was not provided as a cmd line flag, try to read it from
281275
// the specified profile.
282276
isAccountHost := (&config.Config{Host: authArguments.Host, Experimental_IsUnifiedHost: authArguments.IsUnifiedHost}).HostType() != config.WorkspaceHost

‎cmd/auth/login_test.go‎

Lines changed: 8 additions & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -26,8 +26,7 @@ func TestSetHostDoesNotFailWithNoDatabrickscfg(t *testing.T) {
2626
existingProfile, err := loadProfileByName(ctx, "foo", profile.DefaultProfiler)
2727
assert.NoError(t, err)
2828

29-
cmd := newLoginCommand(&auth.AuthArguments{})
30-
err = setHostAndAccountId(ctx, cmd, existingProfile, &auth.AuthArguments{Host: "test"}, []string{})
29+
err = setHostAndAccountId(ctx, existingProfile, &auth.AuthArguments{Host: "test"}, []string{})
3130
assert.NoError(t, err)
3231
}
3332

@@ -39,34 +38,32 @@ func TestSetHost(t *testing.T) {
3938
profile1 := loadTestProfile(t, ctx, "profile-1")
4039
profile2 := loadTestProfile(t, ctx, "profile-2")
4140

42-
cmd := newLoginCommand(&auth.AuthArguments{})
43-
4441
// Test error when both flag and argument are provided
4542
authArguments.Host = "val from --host"
46-
err := setHostAndAccountId(ctx, cmd, profile1, &authArguments, []string{"val from [HOST]"})
43+
err := setHostAndAccountId(ctx, profile1, &authArguments, []string{"val from [HOST]"})
4744
assert.EqualError(t, err, "please only provide a host as an argument or a flag, not both")
4845

4946
// Test setting host from flag
5047
authArguments.Host = "val from --host"
51-
err = setHostAndAccountId(ctx, cmd, profile1, &authArguments, []string{})
48+
err = setHostAndAccountId(ctx, profile1, &authArguments, []string{})
5249
assert.NoError(t, err)
5350
assert.Equal(t, "val from --host", authArguments.Host)
5451

5552
// Test setting host from argument
5653
authArguments.Host = ""
57-
err = setHostAndAccountId(ctx, cmd, profile1, &authArguments, []string{"val from [HOST]"})
54+
err = setHostAndAccountId(ctx, profile1, &authArguments, []string{"val from [HOST]"})
5855
assert.NoError(t, err)
5956
assert.Equal(t, "val from [HOST]", authArguments.Host)
6057

6158
// Test setting host from profile
6259
authArguments.Host = ""
63-
err = setHostAndAccountId(ctx, cmd, profile1, &authArguments, []string{})
60+
err = setHostAndAccountId(ctx, profile1, &authArguments, []string{})
6461
assert.NoError(t, err)
6562
assert.Equal(t, "https://www.host1.com", authArguments.Host)
6663

6764
// Test setting host from profile
6865
authArguments.Host = ""
69-
err = setHostAndAccountId(ctx, cmd, profile2, &authArguments, []string{})
66+
err = setHostAndAccountId(ctx, profile2, &authArguments, []string{})
7067
assert.NoError(t, err)
7168
assert.Equal(t, "https://www.host2.com", authArguments.Host)
7269

@@ -83,18 +80,16 @@ func TestSetAccountId(t *testing.T) {
8380

8481
accountProfile := loadTestProfile(t, ctx, "account-profile")
8582

86-
cmd := newLoginCommand(&auth.AuthArguments{})
87-
8883
// Test setting account-id from flag
8984
authArguments.AccountID = "val from --account-id"
90-
err := setHostAndAccountId(ctx, cmd, accountProfile, &authArguments, []string{})
85+
err := setHostAndAccountId(ctx, accountProfile, &authArguments, []string{})
9186
assert.NoError(t, err)
9287
assert.Equal(t, "https://accounts.cloud.databricks.com", authArguments.Host)
9388
assert.Equal(t, "val from --account-id", authArguments.AccountID)
9489

9590
// Test setting account_id from profile
9691
authArguments.AccountID = ""
97-
err = setHostAndAccountId(ctx, cmd, accountProfile, &authArguments, []string{})
92+
err = setHostAndAccountId(ctx, accountProfile, &authArguments, []string{})
9893
require.NoError(t, err)
9994
assert.Equal(t, "https://accounts.cloud.databricks.com", authArguments.Host)
10095
assert.Equal(t, "id-from-profile", authArguments.AccountID)

‎cmd/auth/token.go‎

Lines changed: 7 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -102,7 +102,12 @@ func loadToken(ctx context.Context, args loadTokenArgs) (*oauth2.Token, error) {
102102
return nil, err
103103
}
104104

105-
err = setHostAndAccountId(ctx, args.cmd, existingProfile, args.authArguments, args.args)
105+
// Set IsUnifiedHost from the profile if the flag wasn't explicitly set.
106+
if !args.cmd.Flag("experimental-is-unified-host").Changed && existingProfile != nil {
107+
args.authArguments.IsUnifiedHost = existingProfile.Experimental_IsUnifiedHost
108+
}
109+
110+
err = setHostAndAccountId(ctx, existingProfile, args.authArguments, args.args)
106111
if err != nil {
107112
return nil, err
108113
}
@@ -133,7 +138,7 @@ func loadToken(ctx context.Context, args loadTokenArgs) (*oauth2.Token, error) {
133138
// This is captured in an acceptance test under "cmd/auth/token".
134139
err = errors.New("cache: databricks OAuth is not configured for this host")
135140
}
136-
if rewritten, rewrittenErr := auth.RewriteAuthError(ctx, args.authArguments.Host, args.authArguments.AccountID, args.profileName, err); rewritten {
141+
if rewritten, rewrittenErr := auth.RewriteAuthError(ctx, args.authArguments.Host, args.authArguments.AccountID, args.profileName, args.authArguments.IsUnifiedHost, err); rewritten {
137142
return nil, rewrittenErr
138143
}
139144
helpMsg := helpfulError(ctx, args.profileName, oauthArgument)

‎cmd/root/auth.go‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -314,7 +314,7 @@ func emptyHttpRequest(ctx context.Context) *http.Request {
314314
}
315315

316316
func renderError(ctx context.Context, cfg *config.Config, err error) error {
317-
if rewritten, newErr := auth.RewriteAuthError(ctx, cfg.Host, cfg.AccountID, cfg.Profile, err); rewritten {
317+
if rewritten, newErr := auth.RewriteAuthError(ctx, cfg.Host, cfg.AccountID, cfg.Profile, cfg.Experimental_IsUnifiedHost, err); rewritten {
318318
return newErr
319319
}
320320
return err

‎libs/auth/error.go‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -10,10 +10,10 @@ import (
1010

1111
// RewriteAuthError rewrites the error message for invalid refresh token error.
1212
// It returns whether the error was rewritten and the rewritten error.
13-
func RewriteAuthError(ctx context.Context, host, accountId, profile string, err error) (bool, error) {
13+
func RewriteAuthError(ctx context.Context, host, accountId, profile string, isUnifiedHost bool, err error) (bool, error) {
1414
target := &u2m.InvalidRefreshTokenError{}
1515
if errors.As(err, &target) {
16-
oauthArgument, err := AuthArguments{host, accountId, false}.ToOAuthArgument() // TODO: pass the IsUnifiedHost flag
16+
oauthArgument, err := AuthArguments{Host: host, AccountID: accountId, IsUnifiedHost: isUnifiedHost}.ToOAuthArgument()
1717
if err != nil {
1818
return false, err
1919
}

0 commit comments

Comments
 (0)