Skip to content

Commit 9b8f930

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

6 files changed

Lines changed: 130 additions & 87 deletions

File tree

‎cmd/auth/login.go‎

Lines changed: 36 additions & 44 deletions
Original file line numberDiff line numberDiff line change
@@ -107,10 +107,23 @@ depends on the existing profiles you have set in your configuration file
107107
}
108108

109109
// Set the host and account-id based on the provided arguments and flags.
110-
err := setHostAndAccountId(ctx, profile.DefaultProfiler, profileName, authArguments, args)
110+
// Load the profile first to get cluster ID and other profile data
111+
existingProfile, err := loadProfileByName(ctx, profileName, profile.DefaultProfiler)
111112
if err != nil {
112113
return err
113114
}
115+
116+
err = setHostAndAccountId(ctx, existingProfile, authArguments, args)
117+
if err != nil {
118+
return err
119+
}
120+
121+
// Extract cluster ID from existing profile if available
122+
clusterID := ""
123+
if existingProfile != nil {
124+
clusterID = existingProfile.ClusterID
125+
}
126+
114127
oauthArgument, err := authArguments.ToOAuthArgument()
115128
if err != nil {
116129
return err
@@ -127,6 +140,7 @@ depends on the existing profiles you have set in your configuration file
127140
Host: authArguments.Host,
128141
AccountID: authArguments.AccountID,
129142
AuthType: "databricks-cli",
143+
ClusterID: clusterID,
130144
}
131145

132146
ctx, cancel := context.WithTimeout(ctx, loginTimeout)
@@ -136,14 +150,6 @@ depends on the existing profiles you have set in your configuration file
136150
return err
137151
}
138152

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-
}
147153
if configureCluster {
148154
w, err := databricks.NewWorkspaceClient((*databricks.Config)(&cfg))
149155
if err != nil {
@@ -190,27 +196,21 @@ depends on the existing profiles you have set in your configuration file
190196
// 1. --account-id flag.
191197
// 2. account-id from the specified profile, if available.
192198
// 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 {
199+
func setHostAndAccountId(ctx context.Context, existingProfile *profile.Profile, authArguments *auth.AuthArguments, args []string) error {
194200
// If both [HOST] and --host are provided, return an error.
195201
host := authArguments.Host
196202
if len(args) > 0 && host != "" {
197203
return errors.New("please only provide a host as an argument or a flag, not both")
198204
}
199205

200206
// 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-
207207
if host == "" {
208208
if len(args) > 0 {
209209
// If [HOST] is provided, set the host to the provided positional argument.
210210
authArguments.Host = args[0]
211-
} else if len(profiles) > 0 && profiles[0].Host != "" {
211+
} else if existingProfile != nil && existingProfile.Host != "" {
212212
// If neither [HOST] nor --host are provided, and the profile has a host, use it.
213-
authArguments.Host = profiles[0].Host
213+
authArguments.Host = existingProfile.Host
214214
} else {
215215
// If neither [HOST] nor --host are provided, and the profile does not have a host,
216216
// then prompt the user for a host.
@@ -227,8 +227,8 @@ func setHostAndAccountId(ctx context.Context, profiler profile.Profiler, profile
227227
isAccountClient := (&config.Config{Host: authArguments.Host}).IsAccountClient()
228228
accountID := authArguments.AccountID
229229
if isAccountClient && accountID == "" {
230-
if len(profiles) > 0 && profiles[0].AccountID != "" {
231-
authArguments.AccountID = profiles[0].AccountID
230+
if existingProfile != nil && existingProfile.AccountID != "" {
231+
authArguments.AccountID = existingProfile.AccountID
232232
} else {
233233
// Prompt user for the account-id if it we could not get it from a
234234
// profile.
@@ -254,32 +254,24 @@ func getProfileName(authArguments *auth.AuthArguments) string {
254254
return split[0]
255255
}
256256

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)
257+
func loadProfileByName(ctx context.Context, profileName string, profiler profile.Profiler) (*profile.Profile, error) {
258+
if profileName == "" {
259+
return nil, nil
265260
}
266261

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
262+
if profiler == nil {
263+
return nil, errors.New("profiler cannot be nil")
282264
}
283265

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

‎cmd/auth/login_test.go‎

Lines changed: 77 additions & 40 deletions
Original file line numberDiff line numberDiff line change
@@ -12,50 +12,65 @@ 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+
// Load profile (will fail but we tolerate it)
27+
existingProfile, err := loadProfileByName(ctx, "foo", profile.DefaultProfiler)
28+
assert.NoError(t, err)
29+
30+
err = setHostAndAccountId(ctx, existingProfile, &auth.AuthArguments{Host: "test"}, []string{})
1931
assert.NoError(t, err)
2032
}
2133

2234
func TestSetHost(t *testing.T) {
23-
authArguments := auth.AuthArguments{}
35+
var authArguments auth.AuthArguments
2436
t.Setenv("DATABRICKS_CONFIG_FILE", "./testdata/.databrickscfg")
2537
ctx, _ := cmdio.SetupTest(context.Background())
2638

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

3247
// Test setting host from flag
3348
authArguments.Host = "val from --host"
34-
err = setHostAndAccountId(ctx, profile.DefaultProfiler, "profile-1", &authArguments, []string{})
49+
err = setHostAndAccountId(ctx, profile1, &authArguments, []string{})
3550
assert.NoError(t, err)
3651
assert.Equal(t, "val from --host", authArguments.Host)
3752

3853
// Test setting host from argument
3954
authArguments.Host = ""
40-
err = setHostAndAccountId(ctx, profile.DefaultProfiler, "profile-1", &authArguments, []string{"val from [HOST]"})
55+
err = setHostAndAccountId(ctx, profile1, &authArguments, []string{"val from [HOST]"})
4156
assert.NoError(t, err)
4257
assert.Equal(t, "val from [HOST]", authArguments.Host)
4358

4459
// Test setting host from profile
4560
authArguments.Host = ""
46-
err = setHostAndAccountId(ctx, profile.DefaultProfiler, "profile-1", &authArguments, []string{})
61+
err = setHostAndAccountId(ctx, profile1, &authArguments, []string{})
4762
assert.NoError(t, err)
4863
assert.Equal(t, "https://www.host1.com", authArguments.Host)
4964

5065
// Test setting host from profile
5166
authArguments.Host = ""
52-
err = setHostAndAccountId(ctx, profile.DefaultProfiler, "profile-2", &authArguments, []string{})
67+
err = setHostAndAccountId(ctx, profile2, &authArguments, []string{})
5368
assert.NoError(t, err)
5469
assert.Equal(t, "https://www.host2.com", authArguments.Host)
5570

5671
// Test host is not set. Should prompt.
5772
authArguments.Host = ""
58-
err = setHostAndAccountId(ctx, profile.DefaultProfiler, "", &authArguments, []string{})
73+
err = setHostAndAccountId(ctx, nil, &authArguments, []string{})
5974
assert.EqualError(t, err, "the command is being run in a non-interactive environment, please specify a host using --host")
6075
}
6176

@@ -64,75 +79,97 @@ func TestSetAccountId(t *testing.T) {
6479
t.Setenv("DATABRICKS_CONFIG_FILE", "./testdata/.databrickscfg")
6580
ctx, _ := cmdio.SetupTest(context.Background())
6681

82+
accountProfile := loadTestProfile(t, ctx, "account-profile")
83+
6784
// Test setting account-id from flag
6885
authArguments.AccountID = "val from --account-id"
69-
err := setHostAndAccountId(ctx, profile.DefaultProfiler, "account-profile", &authArguments, []string{})
86+
err := setHostAndAccountId(ctx, accountProfile, &authArguments, []string{})
7087
assert.NoError(t, err)
7188
assert.Equal(t, "https://accounts.cloud.databricks.com", authArguments.Host)
7289
assert.Equal(t, "val from --account-id", authArguments.AccountID)
7390

7491
// Test setting account_id from profile
7592
authArguments.AccountID = ""
76-
err = setHostAndAccountId(ctx, profile.DefaultProfiler, "account-profile", &authArguments, []string{})
93+
err = setHostAndAccountId(ctx, accountProfile, &authArguments, []string{})
7794
require.NoError(t, err)
7895
assert.Equal(t, "https://accounts.cloud.databricks.com", authArguments.Host)
7996
assert.Equal(t, "id-from-profile", authArguments.AccountID)
8097

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

88-
func TestLoginGetClusterID(t *testing.T) {
105+
func TestLoadProfileByNameAndClusterID(t *testing.T) {
89106
testCases := []struct {
90-
name string
91-
profile string
92-
configFileEnv string
93-
expected string
107+
name string
108+
profile string
109+
configFileEnv string
110+
expectedHost string
111+
expectedClusterID string
94112
}{
95113
{
96-
name: "cluster profile",
97-
profile: "cluster-profile",
98-
configFileEnv: "./testdata/.databrickscfg",
99-
expected: "cluster-from-config",
114+
name: "cluster profile",
115+
profile: "cluster-profile",
116+
configFileEnv: "./testdata/.databrickscfg",
117+
expectedHost: "https://www.host2.com",
118+
expectedClusterID: "cluster-from-config",
100119
},
101120
{
102-
name: "empty profile",
103-
profile: "no-profile",
104-
configFileEnv: "./testdata/.databrickscfg",
105-
expected: "",
121+
name: "empty profile",
122+
profile: "no-profile",
123+
configFileEnv: "./testdata/.databrickscfg",
124+
expectedHost: "",
125+
expectedClusterID: "",
106126
},
107127
{
108-
name: "account profile",
109-
profile: "account-profile",
110-
configFileEnv: "./testdata/.databrickscfg",
111-
expected: "",
128+
name: "account profile",
129+
profile: "account-profile",
130+
configFileEnv: "./testdata/.databrickscfg",
131+
expectedHost: "https://accounts.cloud.databricks.com",
132+
expectedClusterID: "",
112133
},
113134
{
114-
name: "config doesn't exist",
115-
profile: "any-profile",
116-
configFileEnv: "./nonexistent/.databrickscfg",
117-
expected: "",
135+
name: "config doesn't exist",
136+
profile: "any-profile",
137+
configFileEnv: "./nonexistent/.databrickscfg",
138+
expectedHost: "",
139+
expectedClusterID: "",
118140
},
119141
{
120-
name: "invalid profile (missing host)",
121-
profile: "invalid-profile",
122-
configFileEnv: "./testdata/.databrickscfg",
123-
expected: "",
142+
name: "invalid profile (missing host)",
143+
profile: "invalid-profile",
144+
configFileEnv: "./testdata/.databrickscfg",
145+
expectedHost: "",
146+
expectedClusterID: "",
147+
},
148+
{
149+
name: "no config file specified",
150+
profile: "any-profile",
151+
configFileEnv: "",
152+
expectedHost: "",
153+
expectedClusterID: "",
124154
},
125155
}
126156

127157
for _, tc := range testCases {
128158
t.Run(tc.name, func(t *testing.T) {
129159
t.Setenv("DATABRICKS_CONFIG_FILE", tc.configFileEnv)
130160

131-
clusterID, err := getClusterID(context.Background(), tc.profile)
161+
profile, err := loadProfileByName(context.Background(), tc.profile, profile.DefaultProfiler)
132162
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)
163+
164+
if tc.expectedHost == "" {
165+
assert.Nil(t, profile, "Test case '%s' failed: expected nil profile but got non-nil profile", tc.name)
166+
} else {
167+
assert.NotNil(t, profile, "Test case '%s' failed: expected profile but got nil", tc.name)
168+
assert.Equal(t, tc.expectedHost, profile.Host,
169+
"Test case '%s' failed: expected host '%s', but got '%s'", tc.name, tc.expectedHost, profile.Host)
170+
assert.Equal(t, tc.expectedClusterID, profile.ClusterID,
171+
"Test case '%s' failed: expected cluster ID '%s', but got '%s'", tc.name, tc.expectedClusterID, profile.ClusterID)
172+
}
136173
})
137174
}
138175
}

‎cmd/auth/token.go‎

Lines changed: 7 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -92,7 +92,13 @@ 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+
// Load the profile first to get profile data
96+
existingProfile, err := loadProfileByName(ctx, args.profileName, args.profiler)
97+
if err != nil {
98+
return nil, err
99+
}
100+
101+
err = setHostAndAccountId(ctx, existingProfile, args.authArguments, args.args)
96102
if err != nil {
97103
return nil, err
98104
}

‎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)