Skip to content

Commit b73cadd

Browse files
alyssa-dbalyssa-db
committed
added clusterid to profile
1 parent c9569f2 commit b73cadd

6 files changed

Lines changed: 126 additions & 87 deletions

File tree

‎cmd/auth/login.go‎

Lines changed: 34 additions & 44 deletions
Original file line numberDiff line numberDiff line change
@@ -106,11 +106,22 @@ depends on the existing profiles you have set in your configuration file
106106
}
107107
}
108108

109+
existingProfile, err := loadProfileByName(ctx, profileName, profile.DefaultProfiler)
110+
if err != nil {
111+
return err
112+
}
113+
109114
// Set the host and account-id based on the provided arguments and flags.
110-
err := setHostAndAccountId(ctx, profile.DefaultProfiler, profileName, authArguments, args)
115+
err = setHostAndAccountId(ctx, existingProfile, authArguments, args)
111116
if err != nil {
112117
return err
113118
}
119+
120+
clusterID := ""
121+
if existingProfile != nil {
122+
clusterID = existingProfile.ClusterID
123+
}
124+
114125
oauthArgument, err := authArguments.ToOAuthArgument()
115126
if err != nil {
116127
return err
@@ -127,6 +138,7 @@ depends on the existing profiles you have set in your configuration file
127138
Host: authArguments.Host,
128139
AccountID: authArguments.AccountID,
129140
AuthType: "databricks-cli",
141+
ClusterID: clusterID,
130142
}
131143

132144
ctx, cancel := context.WithTimeout(ctx, loginTimeout)
@@ -136,14 +148,6 @@ depends on the existing profiles you have set in your configuration file
136148
return err
137149
}
138150

139-
// Bug fix: Preserve cluster_id from existing profile during login
140-
// Prior to v0.233.0, the login command would overwrite the entire profile configuration,
141-
// causing loss of cluster_id and other settings. Now we first retrieve the existing
142-
// cluster_id before saving the new configuration to ensure it's preserved.
143-
cfg.ClusterID, err = getClusterID(ctx, profileName)
144-
if err != nil {
145-
return err
146-
}
147151
if configureCluster {
148152
w, err := databricks.NewWorkspaceClient((*databricks.Config)(&cfg))
149153
if err != nil {
@@ -190,27 +194,21 @@ depends on the existing profiles you have set in your configuration file
190194
// 1. --account-id flag.
191195
// 2. account-id from the specified profile, if available.
192196
// 3. Prompt the user for the account-id.
193-
func setHostAndAccountId(ctx context.Context, profiler profile.Profiler, profileName string, authArguments *auth.AuthArguments, args []string) error {
197+
func setHostAndAccountId(ctx context.Context, existingProfile *profile.Profile, authArguments *auth.AuthArguments, args []string) error {
194198
// If both [HOST] and --host are provided, return an error.
195199
host := authArguments.Host
196200
if len(args) > 0 && host != "" {
197201
return errors.New("please only provide a host as an argument or a flag, not both")
198202
}
199203

200204
// If the chosen profile has a hostname and the user hasn't specified a host, infer the host from the profile.
201-
profiles, err := profiler.LoadProfiles(ctx, profile.WithName(profileName))
202-
// Tolerate ErrNoConfiguration here, as we will write out a configuration as part of the login flow.
203-
if err != nil && !errors.Is(err, profile.ErrNoConfiguration) {
204-
return err
205-
}
206-
207205
if host == "" {
208206
if len(args) > 0 {
209207
// If [HOST] is provided, set the host to the provided positional argument.
210208
authArguments.Host = args[0]
211-
} else if len(profiles) > 0 && profiles[0].Host != "" {
209+
} else if existingProfile != nil && existingProfile.Host != "" {
212210
// If neither [HOST] nor --host are provided, and the profile has a host, use it.
213-
authArguments.Host = profiles[0].Host
211+
authArguments.Host = existingProfile.Host
214212
} else {
215213
// If neither [HOST] nor --host are provided, and the profile does not have a host,
216214
// then prompt the user for a host.
@@ -227,8 +225,8 @@ func setHostAndAccountId(ctx context.Context, profiler profile.Profiler, profile
227225
isAccountClient := (&config.Config{Host: authArguments.Host}).IsAccountClient()
228226
accountID := authArguments.AccountID
229227
if isAccountClient && accountID == "" {
230-
if len(profiles) > 0 && profiles[0].AccountID != "" {
231-
authArguments.AccountID = profiles[0].AccountID
228+
if existingProfile != nil && existingProfile.AccountID != "" {
229+
authArguments.AccountID = existingProfile.AccountID
232230
} else {
233231
// Prompt user for the account-id if it we could not get it from a
234232
// profile.
@@ -254,32 +252,24 @@ func getProfileName(authArguments *auth.AuthArguments) string {
254252
return split[0]
255253
}
256254

257-
func getClusterID(ctx context.Context, profileName string) (string, error) {
258-
file, err := profile.DefaultProfiler.Get(ctx)
259-
if err != nil {
260-
// If no configuration file exists, return empty string (no error)
261-
if errors.Is(err, profile.ErrNoConfiguration) {
262-
return "", nil
263-
}
264-
return "", fmt.Errorf("cannot load Databricks config file: %w", err)
255+
func loadProfileByName(ctx context.Context, profileName string, profiler profile.Profiler) (*profile.Profile, error) {
256+
if profileName == "" {
257+
return nil, nil
265258
}
266259

267-
// file.Sections() is used in LoadProfiles() to iterate over profiles in login.
268-
// Even with multiple profiles of the same name in the config file,
269-
// it associates each profile name to at most one section.
270-
// It takes the cluster_id if any profiles of the same name have it.
271-
for _, v := range file.Sections() {
272-
if v.Name() != profileName {
273-
continue
274-
}
275-
all := v.KeysHash()
276-
_, ok := all["host"]
277-
if !ok {
278-
// invalid profile, skip
279-
continue
280-
}
281-
return all["cluster_id"], nil
260+
if profiler == nil {
261+
return nil, errors.New("profiler cannot be nil")
282262
}
283263

284-
return "", nil
264+
profiles, err := profiler.LoadProfiles(ctx, profile.WithName(profileName))
265+
// Tolerate ErrNoConfiguration here, as we will write out a configuration as part of the login flow.
266+
if err != nil && !errors.Is(err, profile.ErrNoConfiguration) {
267+
return nil, err
268+
}
269+
270+
if len(profiles) > 0 {
271+
// LoadProfiles returns only one profile per name, even with multiple profiles in the config file with the same name.
272+
return &profiles[0], nil
273+
}
274+
return nil, nil
285275
}

‎cmd/auth/login_test.go‎

Lines changed: 76 additions & 40 deletions
Original file line numberDiff line numberDiff line change
@@ -12,50 +12,64 @@ import (
1212
"github.com/stretchr/testify/require"
1313
)
1414

15+
func loadTestProfile(t *testing.T, ctx context.Context, profileName string) *profile.Profile {
16+
profile, err := loadProfileByName(ctx, profileName, profile.DefaultProfiler)
17+
require.NoError(t, err)
18+
require.NotNil(t, profile)
19+
return profile
20+
}
21+
1522
func TestSetHostDoesNotFailWithNoDatabrickscfg(t *testing.T) {
1623
ctx := context.Background()
1724
ctx = env.Set(ctx, "DATABRICKS_CONFIG_FILE", "./imaginary-file/databrickscfg")
18-
err := setHostAndAccountId(ctx, profile.DefaultProfiler, "foo", &auth.AuthArguments{Host: "test"}, []string{})
25+
26+
existingProfile, err := loadProfileByName(ctx, "foo", profile.DefaultProfiler)
27+
assert.NoError(t, err)
28+
29+
err = setHostAndAccountId(ctx, existingProfile, &auth.AuthArguments{Host: "test"}, []string{})
1930
assert.NoError(t, err)
2031
}
2132

2233
func TestSetHost(t *testing.T) {
23-
authArguments := auth.AuthArguments{}
34+
var authArguments auth.AuthArguments
2435
t.Setenv("DATABRICKS_CONFIG_FILE", "./testdata/.databrickscfg")
2536
ctx, _ := cmdio.SetupTest(context.Background())
2637

38+
profile1 := loadTestProfile(t, ctx, "profile-1")
39+
profile2 := loadTestProfile(t, ctx, "profile-2")
40+
2741
// Test error when both flag and argument are provided
2842
authArguments.Host = "val from --host"
29-
err := setHostAndAccountId(ctx, profile.DefaultProfiler, "profile-1", &authArguments, []string{"val from [HOST]"})
43+
err := setHostAndAccountId(ctx, profile1, &authArguments, []string{"val from [HOST]"})
3044
assert.EqualError(t, err, "please only provide a host as an argument or a flag, not both")
3145

3246
// Test setting host from flag
3347
authArguments.Host = "val from --host"
34-
err = setHostAndAccountId(ctx, profile.DefaultProfiler, "profile-1", &authArguments, []string{})
48+
err = setHostAndAccountId(ctx, profile1, &authArguments, []string{})
3549
assert.NoError(t, err)
3650
assert.Equal(t, "val from --host", authArguments.Host)
3751

3852
// Test setting host from argument
3953
authArguments.Host = ""
40-
err = setHostAndAccountId(ctx, profile.DefaultProfiler, "profile-1", &authArguments, []string{"val from [HOST]"})
54+
err = setHostAndAccountId(ctx, profile1, &authArguments, []string{"val from [HOST]"})
4155
assert.NoError(t, err)
4256
assert.Equal(t, "val from [HOST]", authArguments.Host)
4357

4458
// Test setting host from profile
4559
authArguments.Host = ""
46-
err = setHostAndAccountId(ctx, profile.DefaultProfiler, "profile-1", &authArguments, []string{})
60+
err = setHostAndAccountId(ctx, profile1, &authArguments, []string{})
4761
assert.NoError(t, err)
4862
assert.Equal(t, "https://www.host1.com", authArguments.Host)
4963

5064
// Test setting host from profile
5165
authArguments.Host = ""
52-
err = setHostAndAccountId(ctx, profile.DefaultProfiler, "profile-2", &authArguments, []string{})
66+
err = setHostAndAccountId(ctx, profile2, &authArguments, []string{})
5367
assert.NoError(t, err)
5468
assert.Equal(t, "https://www.host2.com", authArguments.Host)
5569

5670
// Test host is not set. Should prompt.
5771
authArguments.Host = ""
58-
err = setHostAndAccountId(ctx, profile.DefaultProfiler, "", &authArguments, []string{})
72+
err = setHostAndAccountId(ctx, nil, &authArguments, []string{})
5973
assert.EqualError(t, err, "the command is being run in a non-interactive environment, please specify a host using --host")
6074
}
6175

@@ -64,75 +78,97 @@ func TestSetAccountId(t *testing.T) {
6478
t.Setenv("DATABRICKS_CONFIG_FILE", "./testdata/.databrickscfg")
6579
ctx, _ := cmdio.SetupTest(context.Background())
6680

81+
accountProfile := loadTestProfile(t, ctx, "account-profile")
82+
6783
// Test setting account-id from flag
6884
authArguments.AccountID = "val from --account-id"
69-
err := setHostAndAccountId(ctx, profile.DefaultProfiler, "account-profile", &authArguments, []string{})
85+
err := setHostAndAccountId(ctx, accountProfile, &authArguments, []string{})
7086
assert.NoError(t, err)
7187
assert.Equal(t, "https://accounts.cloud.databricks.com", authArguments.Host)
7288
assert.Equal(t, "val from --account-id", authArguments.AccountID)
7389

7490
// Test setting account_id from profile
7591
authArguments.AccountID = ""
76-
err = setHostAndAccountId(ctx, profile.DefaultProfiler, "account-profile", &authArguments, []string{})
92+
err = setHostAndAccountId(ctx, accountProfile, &authArguments, []string{})
7793
require.NoError(t, err)
7894
assert.Equal(t, "https://accounts.cloud.databricks.com", authArguments.Host)
7995
assert.Equal(t, "id-from-profile", authArguments.AccountID)
8096

8197
// Neither flag nor profile account-id is set, should prompt
8298
authArguments.AccountID = ""
8399
authArguments.Host = "https://accounts.cloud.databricks.com"
84-
err = setHostAndAccountId(ctx, profile.DefaultProfiler, "", &authArguments, []string{})
100+
err = setHostAndAccountId(ctx, nil, &authArguments, []string{})
85101
assert.EqualError(t, err, "the command is being run in a non-interactive environment, please specify an account ID using --account-id")
86102
}
87103

88-
func TestLoginGetClusterID(t *testing.T) {
104+
func TestLoadProfileByNameAndClusterID(t *testing.T) {
89105
testCases := []struct {
90-
name string
91-
profile string
92-
configFileEnv string
93-
expected string
106+
name string
107+
profile string
108+
configFileEnv string
109+
expectedHost string
110+
expectedClusterID string
94111
}{
95112
{
96-
name: "cluster profile",
97-
profile: "cluster-profile",
98-
configFileEnv: "./testdata/.databrickscfg",
99-
expected: "cluster-from-config",
113+
name: "cluster profile",
114+
profile: "cluster-profile",
115+
configFileEnv: "./testdata/.databrickscfg",
116+
expectedHost: "https://www.host2.com",
117+
expectedClusterID: "cluster-from-config",
100118
},
101119
{
102-
name: "empty profile",
103-
profile: "no-profile",
104-
configFileEnv: "./testdata/.databrickscfg",
105-
expected: "",
120+
name: "empty profile",
121+
profile: "no-profile",
122+
configFileEnv: "./testdata/.databrickscfg",
123+
expectedHost: "",
124+
expectedClusterID: "",
106125
},
107126
{
108-
name: "account profile",
109-
profile: "account-profile",
110-
configFileEnv: "./testdata/.databrickscfg",
111-
expected: "",
127+
name: "account profile",
128+
profile: "account-profile",
129+
configFileEnv: "./testdata/.databrickscfg",
130+
expectedHost: "https://accounts.cloud.databricks.com",
131+
expectedClusterID: "",
112132
},
113133
{
114-
name: "config doesn't exist",
115-
profile: "any-profile",
116-
configFileEnv: "./nonexistent/.databrickscfg",
117-
expected: "",
134+
name: "config doesn't exist",
135+
profile: "any-profile",
136+
configFileEnv: "./nonexistent/.databrickscfg",
137+
expectedHost: "",
138+
expectedClusterID: "",
118139
},
119140
{
120-
name: "invalid profile (missing host)",
121-
profile: "invalid-profile",
122-
configFileEnv: "./testdata/.databrickscfg",
123-
expected: "",
141+
name: "invalid profile (missing host)",
142+
profile: "invalid-profile",
143+
configFileEnv: "./testdata/.databrickscfg",
144+
expectedHost: "",
145+
expectedClusterID: "",
146+
},
147+
{
148+
name: "no config file specified",
149+
profile: "any-profile",
150+
configFileEnv: "",
151+
expectedHost: "",
152+
expectedClusterID: "",
124153
},
125154
}
126155

127156
for _, tc := range testCases {
128157
t.Run(tc.name, func(t *testing.T) {
129158
t.Setenv("DATABRICKS_CONFIG_FILE", tc.configFileEnv)
130159

131-
clusterID, err := getClusterID(context.Background(), tc.profile)
160+
profile, err := loadProfileByName(context.Background(), tc.profile, profile.DefaultProfiler)
132161
require.NoError(t, err)
133-
assert.Equal(t, tc.expected, clusterID,
134-
"Test case '%s' failed: expected cluster ID '%s', but got '%s' (profile: %s, config file: %s)",
135-
tc.name, tc.expected, clusterID, tc.profile, tc.configFileEnv)
162+
163+
if tc.expectedHost == "" {
164+
assert.Nil(t, profile, "Test case '%s' failed: expected nil profile but got non-nil profile", tc.name)
165+
} else {
166+
assert.NotNil(t, profile, "Test case '%s' failed: expected profile but got nil", tc.name)
167+
assert.Equal(t, tc.expectedHost, profile.Host,
168+
"Test case '%s' failed: expected host '%s', but got '%s'", tc.name, tc.expectedHost, profile.Host)
169+
assert.Equal(t, tc.expectedClusterID, profile.ClusterID,
170+
"Test case '%s' failed: expected cluster ID '%s', but got '%s'", tc.name, tc.expectedClusterID, profile.ClusterID)
171+
}
136172
})
137173
}
138174
}

‎cmd/auth/token.go‎

Lines changed: 6 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -92,7 +92,12 @@ func loadToken(ctx context.Context, args loadTokenArgs) (*oauth2.Token, error) {
9292
return nil, errors.New("providing both a profile and host is not supported")
9393
}
9494

95-
err := setHostAndAccountId(ctx, args.profiler, args.profileName, args.authArguments, args.args)
95+
existingProfile, err := loadProfileByName(ctx, args.profileName, args.profiler)
96+
if err != nil {
97+
return nil, err
98+
}
99+
100+
err = setHostAndAccountId(ctx, existingProfile, args.authArguments, args.args)
96101
if err != nil {
97102
return nil, err
98103
}

‎libs/databrickscfg/profile/file.go‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -82,6 +82,7 @@ func (f FileProfilerImpl) LoadProfiles(ctx context.Context, fn ProfileMatchFunct
8282
Name: v.Name(),
8383
Host: host,
8484
AccountID: all["account_id"],
85+
ClusterID: all["cluster_id"],
8586
}
8687
if fn(profile) {
8788
profiles = append(profiles, profile)

‎libs/databrickscfg/profile/file_test.go‎

Lines changed: 6 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -19,13 +19,18 @@ func TestProfileCloud(t *testing.T) {
1919

2020
func TestProfilesSearchCaseInsensitive(t *testing.T) {
2121
profiles := Profiles{
22-
Profile{Name: "foo", Host: "bar"},
22+
Profile{Name: "foo", Host: "bar", ClusterID: "cluster123"},
2323
}
24+
2425
assert.True(t, profiles.SearchCaseInsensitive("f", 0))
2526
assert.True(t, profiles.SearchCaseInsensitive("OO", 0))
2627
assert.True(t, profiles.SearchCaseInsensitive("b", 0))
2728
assert.True(t, profiles.SearchCaseInsensitive("AR", 0))
2829
assert.False(t, profiles.SearchCaseInsensitive("qu", 0))
30+
assert.True(t, profiles.SearchCaseInsensitive("cluster", 0))
31+
assert.True(t, profiles.SearchCaseInsensitive("123", 0))
32+
assert.True(t, profiles.SearchCaseInsensitive("CLUSTER", 0))
33+
assert.False(t, profiles.SearchCaseInsensitive("xyz", 0))
2934
}
3035

3136
func TestLoadProfilesReturnsHomedirAsTilde(t *testing.T) {

0 commit comments

Comments
 (0)