diff --git a/NEXT_CHANGELOG.md b/NEXT_CHANGELOG.md index ce04d1b9f5..3dd29f1723 100644 --- a/NEXT_CHANGELOG.md +++ b/NEXT_CHANGELOG.md @@ -7,6 +7,7 @@ ### Dependency updates ### CLI +* Fixed an issue where running `databricks auth login` would remove the `cluster_id` field from profiles in `.databrickscfg`. The login process now preserves the `cluster_id` field. ([#2988](https://github.com/databricks/cli/pull/2988)) ### Bundles diff --git a/cmd/auth/login.go b/cmd/auth/login.go index 9fd1ee519e..46f279489a 100644 --- a/cmd/auth/login.go +++ b/cmd/auth/login.go @@ -106,11 +106,22 @@ depends on the existing profiles you have set in your configuration file } } + existingProfile, err := loadProfileByName(ctx, profileName, profile.DefaultProfiler) + if err != nil { + return err + } + // Set the host and account-id based on the provided arguments and flags. - err := setHostAndAccountId(ctx, profile.DefaultProfiler, profileName, authArguments, args) + err = setHostAndAccountId(ctx, existingProfile, authArguments, args) if err != nil { return err } + + clusterID := "" + if existingProfile != nil { + clusterID = existingProfile.ClusterID + } + oauthArgument, err := authArguments.ToOAuthArgument() if err != nil { return err @@ -127,6 +138,7 @@ depends on the existing profiles you have set in your configuration file Host: authArguments.Host, AccountID: authArguments.AccountID, AuthType: "databricks-cli", + ClusterID: clusterID, } ctx, cancel := context.WithTimeout(ctx, loginTimeout) @@ -182,7 +194,7 @@ depends on the existing profiles you have set in your configuration file // 1. --account-id flag. // 2. account-id from the specified profile, if available. // 3. Prompt the user for the account-id. -func setHostAndAccountId(ctx context.Context, profiler profile.Profiler, profileName string, authArguments *auth.AuthArguments, args []string) error { +func setHostAndAccountId(ctx context.Context, existingProfile *profile.Profile, authArguments *auth.AuthArguments, args []string) error { // If both [HOST] and --host are provided, return an error. host := authArguments.Host if len(args) > 0 && host != "" { @@ -190,19 +202,13 @@ func setHostAndAccountId(ctx context.Context, profiler profile.Profiler, profile } // If the chosen profile has a hostname and the user hasn't specified a host, infer the host from the profile. - profiles, err := profiler.LoadProfiles(ctx, profile.WithName(profileName)) - // Tolerate ErrNoConfiguration here, as we will write out a configuration as part of the login flow. - if err != nil && !errors.Is(err, profile.ErrNoConfiguration) { - return err - } - if host == "" { if len(args) > 0 { // If [HOST] is provided, set the host to the provided positional argument. authArguments.Host = args[0] - } else if len(profiles) > 0 && profiles[0].Host != "" { + } else if existingProfile != nil && existingProfile.Host != "" { // If neither [HOST] nor --host are provided, and the profile has a host, use it. - authArguments.Host = profiles[0].Host + authArguments.Host = existingProfile.Host } else { // If neither [HOST] nor --host are provided, and the profile does not have a host, // then prompt the user for a host. @@ -219,8 +225,8 @@ func setHostAndAccountId(ctx context.Context, profiler profile.Profiler, profile isAccountClient := (&config.Config{Host: authArguments.Host}).IsAccountClient() accountID := authArguments.AccountID if isAccountClient && accountID == "" { - if len(profiles) > 0 && profiles[0].AccountID != "" { - authArguments.AccountID = profiles[0].AccountID + if existingProfile != nil && existingProfile.AccountID != "" { + authArguments.AccountID = existingProfile.AccountID } else { // Prompt user for the account-id if it we could not get it from a // profile. @@ -245,3 +251,25 @@ func getProfileName(authArguments *auth.AuthArguments) string { split := strings.Split(host, ".") return split[0] } + +func loadProfileByName(ctx context.Context, profileName string, profiler profile.Profiler) (*profile.Profile, error) { + if profileName == "" { + return nil, nil + } + + if profiler == nil { + return nil, errors.New("profiler cannot be nil") + } + + profiles, err := profiler.LoadProfiles(ctx, profile.WithName(profileName)) + // Tolerate ErrNoConfiguration here, as we will write out a configuration as part of the login flow. + if err != nil && !errors.Is(err, profile.ErrNoConfiguration) { + return nil, err + } + + if len(profiles) > 0 { + // LoadProfiles returns only one profile per name, even with multiple profiles in the config file with the same name. + return &profiles[0], nil + } + return nil, nil +} diff --git a/cmd/auth/login_test.go b/cmd/auth/login_test.go index 8412dab2ce..9b9b8cc526 100644 --- a/cmd/auth/login_test.go +++ b/cmd/auth/login_test.go @@ -12,50 +12,64 @@ import ( "github.com/stretchr/testify/require" ) +func loadTestProfile(t *testing.T, ctx context.Context, profileName string) *profile.Profile { + profile, err := loadProfileByName(ctx, profileName, profile.DefaultProfiler) + require.NoError(t, err) + require.NotNil(t, profile) + return profile +} + func TestSetHostDoesNotFailWithNoDatabrickscfg(t *testing.T) { ctx := context.Background() ctx = env.Set(ctx, "DATABRICKS_CONFIG_FILE", "./imaginary-file/databrickscfg") - err := setHostAndAccountId(ctx, profile.DefaultProfiler, "foo", &auth.AuthArguments{Host: "test"}, []string{}) + + existingProfile, err := loadProfileByName(ctx, "foo", profile.DefaultProfiler) + assert.NoError(t, err) + + err = setHostAndAccountId(ctx, existingProfile, &auth.AuthArguments{Host: "test"}, []string{}) assert.NoError(t, err) } func TestSetHost(t *testing.T) { - authArguments := auth.AuthArguments{} + var authArguments auth.AuthArguments t.Setenv("DATABRICKS_CONFIG_FILE", "./testdata/.databrickscfg") ctx, _ := cmdio.SetupTest(context.Background()) + profile1 := loadTestProfile(t, ctx, "profile-1") + profile2 := loadTestProfile(t, ctx, "profile-2") + // Test error when both flag and argument are provided authArguments.Host = "val from --host" - err := setHostAndAccountId(ctx, profile.DefaultProfiler, "profile-1", &authArguments, []string{"val from [HOST]"}) + err := setHostAndAccountId(ctx, profile1, &authArguments, []string{"val from [HOST]"}) assert.EqualError(t, err, "please only provide a host as an argument or a flag, not both") // Test setting host from flag authArguments.Host = "val from --host" - err = setHostAndAccountId(ctx, profile.DefaultProfiler, "profile-1", &authArguments, []string{}) + err = setHostAndAccountId(ctx, profile1, &authArguments, []string{}) assert.NoError(t, err) assert.Equal(t, "val from --host", authArguments.Host) // Test setting host from argument authArguments.Host = "" - err = setHostAndAccountId(ctx, profile.DefaultProfiler, "profile-1", &authArguments, []string{"val from [HOST]"}) + err = setHostAndAccountId(ctx, profile1, &authArguments, []string{"val from [HOST]"}) assert.NoError(t, err) assert.Equal(t, "val from [HOST]", authArguments.Host) // Test setting host from profile authArguments.Host = "" - err = setHostAndAccountId(ctx, profile.DefaultProfiler, "profile-1", &authArguments, []string{}) + err = setHostAndAccountId(ctx, profile1, &authArguments, []string{}) assert.NoError(t, err) assert.Equal(t, "https://www.host1.com", authArguments.Host) // Test setting host from profile authArguments.Host = "" - err = setHostAndAccountId(ctx, profile.DefaultProfiler, "profile-2", &authArguments, []string{}) + err = setHostAndAccountId(ctx, profile2, &authArguments, []string{}) assert.NoError(t, err) assert.Equal(t, "https://www.host2.com", authArguments.Host) // Test host is not set. Should prompt. authArguments.Host = "" - err = setHostAndAccountId(ctx, profile.DefaultProfiler, "", &authArguments, []string{}) + err = setHostAndAccountId(ctx, nil, &authArguments, []string{}) assert.EqualError(t, err, "the command is being run in a non-interactive environment, please specify a host using --host") } @@ -64,16 +78,18 @@ func TestSetAccountId(t *testing.T) { t.Setenv("DATABRICKS_CONFIG_FILE", "./testdata/.databrickscfg") ctx, _ := cmdio.SetupTest(context.Background()) + accountProfile := loadTestProfile(t, ctx, "account-profile") + // Test setting account-id from flag authArguments.AccountID = "val from --account-id" - err := setHostAndAccountId(ctx, profile.DefaultProfiler, "account-profile", &authArguments, []string{}) + err := setHostAndAccountId(ctx, accountProfile, &authArguments, []string{}) assert.NoError(t, err) assert.Equal(t, "https://accounts.cloud.databricks.com", authArguments.Host) assert.Equal(t, "val from --account-id", authArguments.AccountID) // Test setting account_id from profile authArguments.AccountID = "" - err = setHostAndAccountId(ctx, profile.DefaultProfiler, "account-profile", &authArguments, []string{}) + err = setHostAndAccountId(ctx, accountProfile, &authArguments, []string{}) require.NoError(t, err) assert.Equal(t, "https://accounts.cloud.databricks.com", authArguments.Host) assert.Equal(t, "id-from-profile", authArguments.AccountID) @@ -81,6 +97,93 @@ func TestSetAccountId(t *testing.T) { // Neither flag nor profile account-id is set, should prompt authArguments.AccountID = "" authArguments.Host = "https://accounts.cloud.databricks.com" - err = setHostAndAccountId(ctx, profile.DefaultProfiler, "", &authArguments, []string{}) + err = setHostAndAccountId(ctx, nil, &authArguments, []string{}) assert.EqualError(t, err, "the command is being run in a non-interactive environment, please specify an account ID using --account-id") } + +func TestLoadProfileByNameAndClusterID(t *testing.T) { + testCases := []struct { + name string + profile string + configFileEnv string + homeDirOverride string + expectedHost string + expectedClusterID string + }{ + { + name: "cluster profile", + profile: "cluster-profile", + configFileEnv: "./testdata/.databrickscfg", + expectedHost: "https://www.host2.com", + expectedClusterID: "cluster-from-config", + }, + { + name: "profile from home directory (existing)", + profile: "cluster-profile", + homeDirOverride: "testdata", + expectedHost: "https://www.host2.com", + expectedClusterID: "cluster-from-config", + }, + { + name: "profile does not exist", + profile: "no-profile", + configFileEnv: "./testdata/.databrickscfg", + expectedHost: "", + expectedClusterID: "", + }, + { + name: "account profile", + profile: "account-profile", + configFileEnv: "./testdata/.databrickscfg", + expectedHost: "https://accounts.cloud.databricks.com", + expectedClusterID: "", + }, + { + name: "config doesn't exist", + profile: "any-profile", + configFileEnv: "./nonexistent/.databrickscfg", + expectedHost: "", + expectedClusterID: "", + }, + { + name: "profile from home directory (non-existent)", + profile: "any-profile", + homeDirOverride: "nonexistent", + expectedHost: "", + expectedClusterID: "", + }, + { + name: "invalid profile (missing host)", + profile: "invalid-profile", + configFileEnv: "./testdata/.databrickscfg", + expectedHost: "", + expectedClusterID: "", + }, + } + + for _, tc := range testCases { + t.Run(tc.name, func(t *testing.T) { + ctx := context.Background() + + if tc.configFileEnv != "" { + t.Setenv("DATABRICKS_CONFIG_FILE", tc.configFileEnv) + } else if tc.homeDirOverride != "" { + // Use default ~/.databrickscfg + ctx = env.WithUserHomeDir(ctx, tc.homeDirOverride) + } + + profile, err := loadProfileByName(ctx, tc.profile, profile.DefaultProfiler) + require.NoError(t, err) + + if tc.expectedHost == "" { + assert.Nil(t, profile, "Test case '%s' failed: expected nil profile but got non-nil profile", tc.name) + } else { + assert.NotNil(t, profile, "Test case '%s' failed: expected profile but got nil", tc.name) + assert.Equal(t, tc.expectedHost, profile.Host, + "Test case '%s' failed: expected host '%s', but got '%s'", tc.name, tc.expectedHost, profile.Host) + assert.Equal(t, tc.expectedClusterID, profile.ClusterID, + "Test case '%s' failed: expected cluster ID '%s', but got '%s'", tc.name, tc.expectedClusterID, profile.ClusterID) + } + }) + } +} diff --git a/cmd/auth/testdata/.databrickscfg b/cmd/auth/testdata/.databrickscfg index 06e55224a1..192839b9be 100644 --- a/cmd/auth/testdata/.databrickscfg +++ b/cmd/auth/testdata/.databrickscfg @@ -7,3 +7,11 @@ host = https://www.host2.com [account-profile] host = https://accounts.cloud.databricks.com account_id = id-from-profile + +[cluster-profile] +host = https://www.host2.com +cluster_id = cluster-from-config + +[invalid-profile] +# This profile is missing the required 'host' field +cluster_id = some-cluster-id diff --git a/cmd/auth/token.go b/cmd/auth/token.go index dd0f95732e..5c9e549f77 100644 --- a/cmd/auth/token.go +++ b/cmd/auth/token.go @@ -92,7 +92,12 @@ func loadToken(ctx context.Context, args loadTokenArgs) (*oauth2.Token, error) { return nil, errors.New("providing both a profile and host is not supported") } - err := setHostAndAccountId(ctx, args.profiler, args.profileName, args.authArguments, args.args) + existingProfile, err := loadProfileByName(ctx, args.profileName, args.profiler) + if err != nil { + return nil, err + } + + err = setHostAndAccountId(ctx, existingProfile, args.authArguments, args.args) if err != nil { return nil, err } diff --git a/libs/databrickscfg/profile/file.go b/libs/databrickscfg/profile/file.go index 1b743014e3..9817f53abc 100644 --- a/libs/databrickscfg/profile/file.go +++ b/libs/databrickscfg/profile/file.go @@ -82,6 +82,7 @@ func (f FileProfilerImpl) LoadProfiles(ctx context.Context, fn ProfileMatchFunct Name: v.Name(), Host: host, AccountID: all["account_id"], + ClusterID: all["cluster_id"], } if fn(profile) { profiles = append(profiles, profile) diff --git a/libs/databrickscfg/profile/profile.go b/libs/databrickscfg/profile/profile.go index 510e5c9e55..923f7482a0 100644 --- a/libs/databrickscfg/profile/profile.go +++ b/libs/databrickscfg/profile/profile.go @@ -13,6 +13,7 @@ type Profile struct { Name string Host string AccountID string + ClusterID string } func (p Profile) Cloud() string {