Preserve custom profile fields during 'databricks auth login' (#1897)

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
 

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,27 +194,21 @@ 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 != "" {
 		return errors.New("please only provide a host as an argument or a flag, not both")
 	}
 
 	// 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
+}

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,23 +78,112 @@ 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)
 
 	// 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)
+			}
+		})
+	}
+}

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

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
 	}

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)

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 {