Skip to content

Commit f86a9d4

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

6 files changed

Lines changed: 129 additions & 86 deletions

File tree

‎cmd/auth/login.go‎

Lines changed: 34 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)
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,22 @@ 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+
// loadProfileByName loads a profile by name from the default profiler.
258+
// Returns nil if no profile is found or if there's an error (except ErrNoConfiguration which is tolerated).
259+
func loadProfileByName(ctx context.Context, profileName string) (*profile.Profile, error) {
260+
if profileName == "" {
261+
return nil, nil
265262
}
266263

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
264+
profiles, err := profile.DefaultProfiler.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
282268
}
283269

284-
return "", nil
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: 78 additions & 39 deletions
Original file line numberDiff line numberDiff line change
@@ -6,7 +6,6 @@ import (
66

77
"github.com/databricks/cli/libs/auth"
88
"github.com/databricks/cli/libs/cmdio"
9-
"github.com/databricks/cli/libs/databrickscfg/profile"
109
"github.com/databricks/cli/libs/env"
1110
"github.com/stretchr/testify/assert"
1211
"github.com/stretchr/testify/require"
@@ -15,7 +14,12 @@ import (
1514
func TestSetHostDoesNotFailWithNoDatabrickscfg(t *testing.T) {
1615
ctx := context.Background()
1716
ctx = env.Set(ctx, "DATABRICKS_CONFIG_FILE", "./imaginary-file/databrickscfg")
18-
err := setHostAndAccountId(ctx, profile.DefaultProfiler, "foo", &auth.AuthArguments{Host: "test"}, []string{})
17+
18+
// Load profile (will fail but we tolerate it)
19+
existingProfile, err := loadProfileByName(ctx, "foo")
20+
assert.NoError(t, err)
21+
22+
err = setHostAndAccountId(ctx, existingProfile, &auth.AuthArguments{Host: "test"}, []string{})
1923
assert.NoError(t, err)
2024
}
2125

@@ -24,38 +28,48 @@ func TestSetHost(t *testing.T) {
2428
t.Setenv("DATABRICKS_CONFIG_FILE", "./testdata/.databrickscfg")
2529
ctx, _ := cmdio.SetupTest(context.Background())
2630

31+
// Load profile-1
32+
profile1, err := loadProfileByName(ctx, "profile-1")
33+
require.NoError(t, err)
34+
require.NotNil(t, profile1)
35+
36+
// Load profile-2
37+
profile2, err := loadProfileByName(ctx, "profile-2")
38+
require.NoError(t, err)
39+
require.NotNil(t, profile2)
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,100 @@ func TestSetAccountId(t *testing.T) {
6478
t.Setenv("DATABRICKS_CONFIG_FILE", "./testdata/.databrickscfg")
6579
ctx, _ := cmdio.SetupTest(context.Background())
6680

81+
// Load account-profile
82+
accountProfile, err := loadProfileByName(ctx, "account-profile")
83+
require.NoError(t, err)
84+
require.NotNil(t, accountProfile)
85+
6786
// Test setting account-id from flag
6887
authArguments.AccountID = "val from --account-id"
69-
err := setHostAndAccountId(ctx, profile.DefaultProfiler, "account-profile", &authArguments, []string{})
88+
err = setHostAndAccountId(ctx, accountProfile, &authArguments, []string{})
7089
assert.NoError(t, err)
7190
assert.Equal(t, "https://accounts.cloud.databricks.com", authArguments.Host)
7291
assert.Equal(t, "val from --account-id", authArguments.AccountID)
7392

7493
// Test setting account_id from profile
7594
authArguments.AccountID = ""
76-
err = setHostAndAccountId(ctx, profile.DefaultProfiler, "account-profile", &authArguments, []string{})
95+
err = setHostAndAccountId(ctx, accountProfile, &authArguments, []string{})
7796
require.NoError(t, err)
7897
assert.Equal(t, "https://accounts.cloud.databricks.com", authArguments.Host)
7998
assert.Equal(t, "id-from-profile", authArguments.AccountID)
8099

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

88107
func TestLoginGetClusterID(t *testing.T) {
89108
testCases := []struct {
90-
name string
91-
profile string
92-
configFileEnv string
93-
expected string
109+
name string
110+
profile string
111+
configFileEnv string
112+
expectedHost string
113+
expectedClusterID string
94114
}{
95115
{
96-
name: "cluster profile",
97-
profile: "cluster-profile",
98-
configFileEnv: "./testdata/.databrickscfg",
99-
expected: "cluster-from-config",
116+
name: "cluster profile",
117+
profile: "cluster-profile",
118+
configFileEnv: "./testdata/.databrickscfg",
119+
expectedHost: "https://www.host2.com",
120+
expectedClusterID: "cluster-from-config",
100121
},
101122
{
102-
name: "empty profile",
103-
profile: "no-profile",
104-
configFileEnv: "./testdata/.databrickscfg",
105-
expected: "",
123+
name: "empty profile",
124+
profile: "no-profile",
125+
configFileEnv: "./testdata/.databrickscfg",
126+
expectedHost: "",
127+
expectedClusterID: "",
106128
},
107129
{
108-
name: "account profile",
109-
profile: "account-profile",
110-
configFileEnv: "./testdata/.databrickscfg",
111-
expected: "",
130+
name: "account profile",
131+
profile: "account-profile",
132+
configFileEnv: "./testdata/.databrickscfg",
133+
expectedHost: "https://accounts.cloud.databricks.com",
134+
expectedClusterID: "",
112135
},
113136
{
114-
name: "config doesn't exist",
115-
profile: "any-profile",
116-
configFileEnv: "./nonexistent/.databrickscfg",
117-
expected: "",
137+
name: "config doesn't exist",
138+
profile: "any-profile",
139+
configFileEnv: "./nonexistent/.databrickscfg",
140+
expectedHost: "",
141+
expectedClusterID: "",
118142
},
119143
{
120-
name: "invalid profile (missing host)",
121-
profile: "invalid-profile",
122-
configFileEnv: "./testdata/.databrickscfg",
123-
expected: "",
144+
name: "invalid profile (missing host)",
145+
profile: "invalid-profile",
146+
configFileEnv: "./testdata/.databrickscfg",
147+
expectedHost: "",
148+
expectedClusterID: "",
149+
},
150+
{
151+
name: "~/.databrickscfg does not exist",
152+
profile: "any-profile",
153+
configFileEnv: "",
154+
expectedHost: "",
155+
expectedClusterID: "",
124156
},
125157
}
126158

127159
for _, tc := range testCases {
128160
t.Run(tc.name, func(t *testing.T) {
129161
t.Setenv("DATABRICKS_CONFIG_FILE", tc.configFileEnv)
130162

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

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