diff --git a/acceptance/cmd/auth/login/nominal/out.databrickscfg b/acceptance/cmd/auth/login/nominal/out.databrickscfg index 99c7d54d1e..75aaea5f41 100644 --- a/acceptance/cmd/auth/login/nominal/out.databrickscfg +++ b/acceptance/cmd/auth/login/nominal/out.databrickscfg @@ -4,3 +4,6 @@ [test] host = [DATABRICKS_URL] auth_type = databricks-cli + +[__databricks-settings__] +default_profile = test diff --git a/acceptance/cmd/auth/login/nominal/output.txt b/acceptance/cmd/auth/login/nominal/output.txt index b42bbd5527..4200636bc2 100644 --- a/acceptance/cmd/auth/login/nominal/output.txt +++ b/acceptance/cmd/auth/login/nominal/output.txt @@ -3,5 +3,5 @@ Profile test was successfully saved >>> [CLI] auth profiles -Name Host Valid -test [DATABRICKS_URL] YES +Name Host Valid +test (Default) [DATABRICKS_URL] YES diff --git a/acceptance/cmd/auth/login/with-scopes/out.databrickscfg b/acceptance/cmd/auth/login/with-scopes/out.databrickscfg index 7aac4e9365..5a3a13b312 100644 --- a/acceptance/cmd/auth/login/with-scopes/out.databrickscfg +++ b/acceptance/cmd/auth/login/with-scopes/out.databrickscfg @@ -5,3 +5,6 @@ host = [DATABRICKS_URL] scopes = jobs,pipelines,clusters auth_type = databricks-cli + +[__databricks-settings__] +default_profile = scoped-test diff --git a/cmd/auth/auth.go b/cmd/auth/auth.go index 4c783fd0e6..2903a676a4 100644 --- a/cmd/auth/auth.go +++ b/cmd/auth/auth.go @@ -33,6 +33,7 @@ GCP: https://docs.gcp.databricks.com/dev-tools/auth/index.html`, cmd.AddCommand(newProfilesCommand()) cmd.AddCommand(newTokenCommand(&authArguments)) cmd.AddCommand(newDescribeCommand()) + cmd.AddCommand(newSwitchCommand()) return cmd } diff --git a/cmd/auth/describe.go b/cmd/auth/describe.go index c21eab376c..4fc832f7e9 100644 --- a/cmd/auth/describe.go +++ b/cmd/auth/describe.go @@ -8,6 +8,7 @@ import ( "github.com/databricks/cli/cmd/root" "github.com/databricks/cli/libs/cmdctx" "github.com/databricks/cli/libs/cmdio" + "github.com/databricks/cli/libs/databrickscfg" "github.com/databricks/cli/libs/flags" "github.com/databricks/databricks-sdk-go/config" "github.com/spf13/cobra" @@ -182,6 +183,9 @@ func getAuthDetails(cmd *cobra.Command, cfg *config.Config, showSensitive bool) profile := cfg.Profile if profile == "" { profile = "default" + if resolved, err := databrickscfg.GetConfiguredDefaultProfile(cmd.Context(), cfg.ConfigFile); err == nil && resolved != "" { + profile = fmt.Sprintf("default (%s)", resolved) + } } details.Configuration["profile"] = &config.AttrConfig{Value: profile, Source: config.Source{Type: config.SourceDynamicConfig}} } diff --git a/cmd/auth/login.go b/cmd/auth/login.go index deab15d286..f11c1e53d8 100644 --- a/cmd/auth/login.go +++ b/cmd/auth/login.go @@ -17,6 +17,7 @@ import ( "github.com/databricks/cli/libs/databrickscfg/profile" "github.com/databricks/cli/libs/env" "github.com/databricks/cli/libs/exec" + "github.com/databricks/cli/libs/log" "github.com/databricks/databricks-sdk-go" "github.com/databricks/databricks-sdk-go/config" "github.com/databricks/databricks-sdk-go/config/experimental/auth/authconv" @@ -241,6 +242,12 @@ depends on the existing profiles you have set in your configuration file } if profileName != "" { + configFile := os.Getenv("DATABRICKS_CONFIG_FILE") + + // Check if this is a brand new profile with no other profiles in the file. + // If so, we'll auto-set it as the default after saving. + isFirstProfile := existingProfile == nil && hasNoProfiles(ctx, profile.DefaultProfiler) + err := databrickscfg.SaveToProfile(ctx, &config.Config{ Profile: profileName, Host: authArguments.Host, @@ -249,7 +256,7 @@ depends on the existing profiles you have set in your configuration file WorkspaceID: authArguments.WorkspaceID, Experimental_IsUnifiedHost: authArguments.IsUnifiedHost, ClusterID: clusterID, - ConfigFile: os.Getenv("DATABRICKS_CONFIG_FILE"), + ConfigFile: configFile, ServerlessComputeID: serverlessComputeID, Scopes: scopesList, }, clearKeys...) @@ -257,6 +264,12 @@ depends on the existing profiles you have set in your configuration file return err } + if isFirstProfile { + if err := databrickscfg.SetDefaultProfile(ctx, profileName, configFile); err != nil { + log.Debugf(ctx, "Failed to auto-set default profile: %v", err) + } + } + cmdio.LogString(ctx, fmt.Sprintf("Profile %s was successfully saved", profileName)) } @@ -415,6 +428,17 @@ func openURLSuppressingStderr(url string) error { return browserpkg.OpenURL(url) } +// hasNoProfiles returns true if the config file has no existing profiles. +// Used to detect first-profile creation so we can auto-set it as default. +// Returns true when the config file doesn't exist yet (ErrNoConfiguration). +func hasNoProfiles(ctx context.Context, profiler profile.Profiler) bool { + profiles, err := profiler.LoadProfiles(ctx, profile.MatchAllProfiles) + if err != nil { + return errors.Is(err, profile.ErrNoConfiguration) + } + return len(profiles) == 0 +} + // oauthLoginClearKeys returns profile keys that should be explicitly removed // when performing an OAuth login. Derives auth credential fields dynamically // from the SDK's ConfigAttributes to stay in sync as new auth methods are added. diff --git a/cmd/auth/login_test.go b/cmd/auth/login_test.go index 013786c4bf..c783a73d4e 100644 --- a/cmd/auth/login_test.go +++ b/cmd/auth/login_test.go @@ -2,6 +2,8 @@ package auth import ( "context" + "os" + "path/filepath" "testing" "github.com/databricks/cli/libs/auth" @@ -255,3 +257,30 @@ func TestLoadProfileByNameAndClusterID(t *testing.T) { }) } } + +func TestHasNoProfiles_FreshMachine(t *testing.T) { + // On a fresh machine there is no config file. LoadProfiles returns + // ErrNoConfiguration. hasNoProfiles must treat this as "no profiles" + // (return true), not as an error (return false). + ctx := t.Context() + t.Setenv("DATABRICKS_CONFIG_FILE", filepath.Join(t.TempDir(), "nonexistent")) + assert.True(t, hasNoProfiles(ctx, profile.DefaultProfiler)) +} + +func TestHasNoProfiles_EmptyFile(t *testing.T) { + ctx := t.Context() + dir := t.TempDir() + configFile := filepath.Join(dir, ".databrickscfg") + require.NoError(t, os.WriteFile(configFile, []byte(""), 0o600)) + t.Setenv("DATABRICKS_CONFIG_FILE", configFile) + assert.True(t, hasNoProfiles(ctx, profile.DefaultProfiler)) +} + +func TestHasNoProfiles_WithExistingProfile(t *testing.T) { + ctx := t.Context() + dir := t.TempDir() + configFile := filepath.Join(dir, ".databrickscfg") + require.NoError(t, os.WriteFile(configFile, []byte("[p1]\nhost = https://abc\n"), 0o600)) + t.Setenv("DATABRICKS_CONFIG_FILE", configFile) + assert.False(t, hasNoProfiles(ctx, profile.DefaultProfiler)) +} diff --git a/cmd/auth/profiles.go b/cmd/auth/profiles.go index 6ee0ba49cc..dd53190842 100644 --- a/cmd/auth/profiles.go +++ b/cmd/auth/profiles.go @@ -10,6 +10,7 @@ import ( "time" "github.com/databricks/cli/libs/cmdio" + "github.com/databricks/cli/libs/databrickscfg" "github.com/databricks/cli/libs/databrickscfg/profile" "github.com/databricks/cli/libs/log" "github.com/databricks/databricks-sdk-go" @@ -26,6 +27,7 @@ type profileMetadata struct { Cloud string `json:"cloud"` AuthType string `json:"auth_type"` Valid bool `json:"valid"` + Default bool `json:"default,omitempty"` } func (c *profileMetadata) IsEmpty() bool { @@ -92,7 +94,7 @@ func newProfilesCommand() *cobra.Command { Annotations: map[string]string{ "template": cmdio.Heredoc(` {{header "Name"}} {{header "Host"}} {{header "Valid"}} - {{range .Profiles}}{{.Name | green}} {{.Host|cyan}} {{bool .Valid}} + {{range .Profiles}}{{.Name | green}}{{if .Default}} (Default){{end}} {{.Host|cyan}} {{bool .Valid}} {{end}}`), }, } @@ -111,6 +113,9 @@ func newProfilesCommand() *cobra.Command { } else if err != nil { return fmt.Errorf("cannot parse config file: %w", err) } + + defaultProfile := databrickscfg.GetDefaultProfileFrom(iniFile) + var wg sync.WaitGroup for _, v := range iniFile.Sections() { hash := v.KeysHash() @@ -119,6 +124,7 @@ func newProfilesCommand() *cobra.Command { Host: hash["host"], AccountID: hash["account_id"], WorkspaceID: hash["workspace_id"], + Default: v.Name() == defaultProfile, } if profile.IsEmpty() { continue diff --git a/cmd/auth/profiles_test.go b/cmd/auth/profiles_test.go index 1dfe662889..afd7b0b548 100644 --- a/cmd/auth/profiles_test.go +++ b/cmd/auth/profiles_test.go @@ -43,3 +43,34 @@ func TestProfiles(t *testing.T) { assert.Equal(t, "aws", profile.Cloud) assert.Equal(t, "pat", profile.AuthType) } + +func TestProfilesDefaultMarker(t *testing.T) { + ctx := t.Context() + dir := t.TempDir() + configFile := filepath.Join(dir, ".databrickscfg") + + // Create two profiles. + for _, name := range []string{"profile-a", "profile-b"} { + err := databrickscfg.SaveToProfile(ctx, &config.Config{ + ConfigFile: configFile, + Profile: name, + Host: "https://" + name + ".cloud.databricks.com", + Token: "token", + }) + require.NoError(t, err) + } + + // Set profile-a as the default. + err := databrickscfg.SetDefaultProfile(ctx, "profile-a", configFile) + require.NoError(t, err) + + t.Setenv("HOME", dir) + if runtime.GOOS == "windows" { + t.Setenv("USERPROFILE", dir) + } + + // Read back the default profile and verify. + defaultProfile, err := databrickscfg.GetDefaultProfile(ctx, configFile) + require.NoError(t, err) + assert.Equal(t, "profile-a", defaultProfile) +} diff --git a/cmd/auth/switch.go b/cmd/auth/switch.go new file mode 100644 index 0000000000..8f32e34aa3 --- /dev/null +++ b/cmd/auth/switch.go @@ -0,0 +1,116 @@ +package auth + +import ( + "context" + "errors" + "fmt" + "os" + "strings" + + "github.com/databricks/cli/libs/cmdio" + "github.com/databricks/cli/libs/databrickscfg" + "github.com/databricks/cli/libs/databrickscfg/profile" + "github.com/manifoldco/promptui" + "github.com/spf13/cobra" +) + +func newSwitchCommand() *cobra.Command { + cmd := &cobra.Command{ + Use: "switch", + Short: "Set the default profile", + Long: `Set a named profile as the default in ~/.databrickscfg. + +The selected profile name is stored in a [__databricks-settings__] section +in the config file under the default_profile key. Use "databricks auth profiles" +to see which profile is currently the default.`, + Args: cobra.NoArgs, + } + + cmd.RunE = func(cmd *cobra.Command, args []string) error { + ctx := cmd.Context() + configFile := os.Getenv("DATABRICKS_CONFIG_FILE") + + profileName := cmd.Flag("profile").Value.String() + + if profileName == "" { + if !cmdio.IsPromptSupported(ctx) { + return errors.New("the command is being run in a non-interactive environment, please specify a profile using --profile") + } + + allProfiles, err := profile.DefaultProfiler.LoadProfiles(ctx, profile.MatchAllProfiles) + if err != nil { + return err + } + if len(allProfiles) == 0 { + return errors.New("no profiles configured. Run 'databricks auth login' to create a profile") + } + + // Use the already-loaded config file to resolve the current default, + // avoiding a redundant file read. + currentDefault := "" + if iniFile, err := profile.DefaultProfiler.Get(ctx); err == nil { + currentDefault = databrickscfg.GetDefaultProfileFrom(iniFile) + } + selectedName, err := promptForSwitchProfile(ctx, allProfiles, currentDefault) + if err != nil { + return err + } + profileName = selectedName + } else { + // Validate the profile exists. + profiles, err := profile.DefaultProfiler.LoadProfiles(ctx, profile.WithName(profileName)) + if err != nil { + return err + } + if len(profiles) == 0 { + return fmt.Errorf("profile %q not found", profileName) + } + } + + err := databrickscfg.SetDefaultProfile(ctx, profileName, configFile) + if err != nil { + return err + } + + cmdio.LogString(ctx, fmt.Sprintf("Default profile set to %q.", profileName)) + return nil + } + + return cmd +} + +// promptForSwitchProfile shows an interactive profile picker for the switch command. +// Reuses profileSelectItem from token.go for consistent display. +func promptForSwitchProfile(ctx context.Context, profiles profile.Profiles, currentDefault string) (string, error) { + items := make([]profileSelectItem, 0, len(profiles)) + for _, p := range profiles { + items = append(items, profileSelectItem{Name: p.Name, Host: p.Host}) + } + + label := "Select a profile to set as default" + if currentDefault != "" { + label = fmt.Sprintf("Current default: %s. Select a new default", currentDefault) + } + + i, _, err := cmdio.RunSelect(ctx, &promptui.Select{ + Label: label, + Items: items, + StartInSearchMode: len(profiles) > 5, + Searcher: func(input string, index int) bool { + input = strings.ToLower(input) + name := strings.ToLower(items[index].Name) + host := strings.ToLower(items[index].Host) + return strings.Contains(name, input) || strings.Contains(host, input) + }, + Templates: &promptui.SelectTemplates{ + Label: "{{ . | faint }}", + Active: `{{.Name | bold}}{{if .Host}} ({{.Host|faint}}){{end}}`, + Inactive: `{{.Name}}{{if .Host}} ({{.Host}}){{end}}`, + Selected: `{{ "Default profile" | faint }}: {{ .Name | bold }}`, + }, + }) + if err != nil { + return "", err + } + return profiles[i].Name, nil +} diff --git a/cmd/auth/switch_test.go b/cmd/auth/switch_test.go new file mode 100644 index 0000000000..74feb057c5 --- /dev/null +++ b/cmd/auth/switch_test.go @@ -0,0 +1,143 @@ +package auth + +import ( + "fmt" + "os" + "path/filepath" + "testing" + + "github.com/databricks/cli/libs/cmdio" + "github.com/databricks/cli/libs/databrickscfg" + "github.com/databricks/databricks-sdk-go/config" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +func TestSwitchCommand_WithProfileFlag(t *testing.T) { + ctx := t.Context() + dir := t.TempDir() + configFile := filepath.Join(dir, ".databrickscfg") + + err := databrickscfg.SaveToProfile(ctx, &config.Config{ + ConfigFile: configFile, + Profile: "my-workspace", + Host: "https://abc.cloud.databricks.com", + Token: "token1", + }) + require.NoError(t, err) + + t.Setenv("DATABRICKS_CONFIG_FILE", configFile) + + ctx = cmdio.MockDiscard(ctx) + + cmd := New() + cmd.PersistentFlags().StringP("profile", "p", "", "~/.databrickscfg profile") + cmd.SetContext(ctx) + cmd.SetArgs([]string{"switch", "--profile", "my-workspace"}) + + err = cmd.Execute() + require.NoError(t, err) + + got, err := databrickscfg.GetDefaultProfile(ctx, configFile) + require.NoError(t, err) + assert.Equal(t, "my-workspace", got) +} + +func TestSwitchCommand_ProfileNotFound(t *testing.T) { + ctx := t.Context() + dir := t.TempDir() + configFile := filepath.Join(dir, ".databrickscfg") + + err := databrickscfg.SaveToProfile(ctx, &config.Config{ + ConfigFile: configFile, + Profile: "my-workspace", + Host: "https://abc.cloud.databricks.com", + Token: "token1", + }) + require.NoError(t, err) + + t.Setenv("DATABRICKS_CONFIG_FILE", configFile) + + ctx = cmdio.MockDiscard(ctx) + + cmd := New() + cmd.PersistentFlags().StringP("profile", "p", "", "~/.databrickscfg profile") + cmd.SetContext(ctx) + cmd.SetArgs([]string{"switch", "--profile", "nonexistent"}) + + err = cmd.Execute() + assert.ErrorContains(t, err, `profile "nonexistent" not found`) +} + +func TestSwitchCommand_NonInteractiveNoProfile(t *testing.T) { + ctx := t.Context() + dir := t.TempDir() + configFile := filepath.Join(dir, ".databrickscfg") + + err := databrickscfg.SaveToProfile(ctx, &config.Config{ + ConfigFile: configFile, + Profile: "my-workspace", + Host: "https://abc.cloud.databricks.com", + Token: "token1", + }) + require.NoError(t, err) + + t.Setenv("DATABRICKS_CONFIG_FILE", configFile) + + ctx = cmdio.MockDiscard(ctx) + + cmd := New() + cmd.PersistentFlags().StringP("profile", "p", "", "~/.databrickscfg profile") + cmd.SetContext(ctx) + cmd.SetArgs([]string{"switch"}) + + err = cmd.Execute() + assert.ErrorContains(t, err, "non-interactive environment") +} + +func TestSwitchCommand_WritesSettingsSection(t *testing.T) { + ctx := t.Context() + dir := t.TempDir() + configFile := filepath.Join(dir, ".databrickscfg") + + for _, name := range []string{"profile-a", "profile-b"} { + err := databrickscfg.SaveToProfile(ctx, &config.Config{ + ConfigFile: configFile, + Profile: name, + Host: fmt.Sprintf("https://%s.cloud.databricks.com", name), + Token: "token", + }) + require.NoError(t, err) + } + + t.Setenv("DATABRICKS_CONFIG_FILE", configFile) + + ctx = cmdio.MockDiscard(ctx) + + cmd := New() + cmd.PersistentFlags().StringP("profile", "p", "", "~/.databrickscfg profile") + cmd.SetContext(ctx) + cmd.SetArgs([]string{"switch", "--profile", "profile-a"}) + + err := cmd.Execute() + require.NoError(t, err) + + // Verify the [__databricks-settings__] section was written. + contents, err := os.ReadFile(configFile) + require.NoError(t, err) + assert.Contains(t, string(contents), "[__databricks-settings__]") + assert.Contains(t, string(contents), "default_profile = profile-a") + + // Switch to another profile. + cmd = New() + cmd.PersistentFlags().StringP("profile", "p", "", "~/.databrickscfg profile") + cmd.SetContext(ctx) + cmd.SetArgs([]string{"switch", "--profile", "profile-b"}) + + err = cmd.Execute() + require.NoError(t, err) + + got, err := databrickscfg.GetDefaultProfile(ctx, configFile) + require.NoError(t, err) + assert.Equal(t, "profile-b", got) +} diff --git a/cmd/auth/token.go b/cmd/auth/token.go index ca8582bd02..13396b9098 100644 --- a/cmd/auth/token.go +++ b/cmd/auth/token.go @@ -14,6 +14,7 @@ import ( "github.com/databricks/cli/libs/databrickscfg" "github.com/databricks/cli/libs/databrickscfg/profile" "github.com/databricks/cli/libs/env" + "github.com/databricks/cli/libs/log" "github.com/databricks/databricks-sdk-go/config" "github.com/databricks/databricks-sdk-go/credentials/u2m" "github.com/databricks/databricks-sdk-go/credentials/u2m/cache" @@ -467,6 +468,10 @@ func runInlineLogin(ctx context.Context, profiler profile.Profiler) (string, *pr if !loginArgs.IsUnifiedHost { clearKeys = append(clearKeys, "experimental_is_unified_host") } + + configFile := os.Getenv("DATABRICKS_CONFIG_FILE") + firstProfile := hasNoProfiles(ctx, profiler) + err = databrickscfg.SaveToProfile(ctx, &config.Config{ Profile: profileName, Host: loginArgs.Host, @@ -474,13 +479,19 @@ func runInlineLogin(ctx context.Context, profiler profile.Profiler) (string, *pr AccountID: loginArgs.AccountID, WorkspaceID: loginArgs.WorkspaceID, Experimental_IsUnifiedHost: loginArgs.IsUnifiedHost, - ConfigFile: os.Getenv("DATABRICKS_CONFIG_FILE"), + ConfigFile: configFile, Scopes: scopesList, }, clearKeys...) if err != nil { return "", nil, err } + if firstProfile { + if err := databrickscfg.SetDefaultProfile(ctx, profileName, configFile); err != nil { + log.Debugf(ctx, "Failed to auto-set default profile: %v", err) + } + } + cmdio.LogString(ctx, fmt.Sprintf("Profile %s was successfully saved", profileName)) p, err := loadProfileByName(ctx, profileName, profiler) diff --git a/cmd/root/auth.go b/cmd/root/auth.go index 60360bd3e6..e292be7974 100644 --- a/cmd/root/auth.go +++ b/cmd/root/auth.go @@ -9,7 +9,9 @@ import ( "github.com/databricks/cli/libs/auth" "github.com/databricks/cli/libs/cmdctx" "github.com/databricks/cli/libs/cmdio" + "github.com/databricks/cli/libs/databrickscfg" "github.com/databricks/cli/libs/databrickscfg/profile" + envlib "github.com/databricks/cli/libs/env" "github.com/databricks/cli/libs/logdiag" "github.com/databricks/databricks-sdk-go" "github.com/databricks/databricks-sdk-go/config" @@ -194,6 +196,17 @@ func MustWorkspaceClient(cmd *cobra.Command, args []string) error { cfg.Profile = profile } + // If --profile and DATABRICKS_CONFIG_PROFILE are both unset, honor the + // explicit [__databricks-settings__].default_profile setting before the + // SDK falls back to the DEFAULT section. + if cfg.Profile == "" && envlib.Get(ctx, "DATABRICKS_CONFIG_PROFILE") == "" { + configFilePath := envlib.Get(ctx, "DATABRICKS_CONFIG_FILE") + resolvedProfile, err := databrickscfg.GetConfiguredDefaultProfile(ctx, configFilePath) + if err == nil && resolvedProfile != "" { + cfg.Profile = resolvedProfile + } + } + _, isTargetFlagSet := targetFlagValue(cmd) // If the profile flag is set but the target flag is not, we should skip loading the bundle configuration. if !isTargetFlagSet && hasProfileFlag { diff --git a/cmd/root/auth_test.go b/cmd/root/auth_test.go index 6e03e5687e..33b7889654 100644 --- a/cmd/root/auth_test.go +++ b/cmd/root/auth_test.go @@ -323,3 +323,114 @@ func TestMustAnyClientWithEmptyDatabricksCfg(t *testing.T) { _, err = MustAnyClient(cmd, []string{}) require.ErrorContains(t, err, "does not contain account profiles") } + +func TestMustWorkspaceClientDefaultProfilePrecedence(t *testing.T) { + testutil.CleanupEnvironment(t) + + configFile := filepath.Join(t.TempDir(), ".databrickscfg") + err := os.WriteFile(configFile, []byte(` +[__databricks-settings__] +default_profile = settings-profile + +[DEFAULT] +host = https://default.cloud.databricks.com +token = default-token + +[settings-profile] +host = https://settings.cloud.databricks.com +token = settings-token + +[env-profile] +host = https://env.cloud.databricks.com +token = env-token + +[flag-profile] +host = https://flag.cloud.databricks.com +token = flag-token +`), 0o600) + require.NoError(t, err) + + testCases := []struct { + name string + profileFlag string + envProfile string + wantProfile string + wantHost string + }{ + { + name: "settings default is used when flag and env are unset", + wantProfile: "settings-profile", + wantHost: "https://settings.cloud.databricks.com", + }, + { + name: "env var takes precedence over settings default", + envProfile: "env-profile", + wantProfile: "env-profile", + wantHost: "https://env.cloud.databricks.com", + }, + { + name: "profile flag takes precedence over env var", + profileFlag: "flag-profile", + envProfile: "env-profile", + wantProfile: "flag-profile", + wantHost: "https://flag.cloud.databricks.com", + }, + } + + for _, tc := range testCases { + t.Run(tc.name, func(t *testing.T) { + testutil.CleanupEnvironment(t) + t.Setenv("DATABRICKS_CONFIG_FILE", configFile) + if tc.envProfile != "" { + t.Setenv("DATABRICKS_CONFIG_PROFILE", tc.envProfile) + } + + ctx := cmdio.MockDiscard(t.Context()) + ctx = SkipLoadBundle(ctx) + cmd := New(ctx) + + if tc.profileFlag != "" { + err := cmd.Flag("profile").Value.Set(tc.profileFlag) + require.NoError(t, err) + } + + err := MustWorkspaceClient(cmd, []string{}) + require.NoError(t, err) + + w := cmdctx.WorkspaceClient(cmd.Context()) + require.NotNil(t, w) + assert.Equal(t, tc.wantProfile, w.Config.Profile) + assert.Equal(t, tc.wantHost, w.Config.Host) + }) + } +} + +func TestMustWorkspaceClientWithoutConfiguredDefaultFallsBackToDefaultSection(t *testing.T) { + testutil.CleanupEnvironment(t) + + configFile := filepath.Join(t.TempDir(), ".databrickscfg") + err := os.WriteFile(configFile, []byte(` +[DEFAULT] +host = https://default.cloud.databricks.com +token = default-token + +[named-profile] +host = https://named.cloud.databricks.com +token = named-token +`), 0o600) + require.NoError(t, err) + + t.Setenv("DATABRICKS_CONFIG_FILE", configFile) + + ctx := cmdio.MockDiscard(t.Context()) + ctx = SkipLoadBundle(ctx) + cmd := New(ctx) + + err = MustWorkspaceClient(cmd, []string{}) + require.NoError(t, err) + + w := cmdctx.WorkspaceClient(cmd.Context()) + require.NotNil(t, w) + assert.Equal(t, "", w.Config.Profile) + assert.Equal(t, "https://default.cloud.databricks.com", w.Config.Host) +} diff --git a/libs/databrickscfg/ops.go b/libs/databrickscfg/ops.go index bf602b6c60..37d89d7ad8 100644 --- a/libs/databrickscfg/ops.go +++ b/libs/databrickscfg/ops.go @@ -18,19 +18,176 @@ const fileMode = 0o600 const defaultComment = "The profile defined in the DEFAULT section is to be used as a fallback when no profile is explicitly specified." -func loadOrCreateConfigFile(ctx context.Context, filename string) (*config.File, error) { +const databricksSettingsSection = "__databricks-settings__" + +// GetConfiguredDefaultProfile returns the explicitly configured default profile +// by loading the config file at configFilePath. +// Returns "" if the file doesn't exist or default_profile is not set. +func GetConfiguredDefaultProfile(ctx context.Context, configFilePath string) (string, error) { + configFile, err := loadConfigFile(ctx, configFilePath) + if err != nil { + return "", err + } + if configFile == nil { + return "", nil + } + return GetConfiguredDefaultProfileFrom(configFile), nil +} + +// GetConfiguredDefaultProfileFrom returns the explicit default profile from +// [__databricks-settings__].default_profile, or "" when it is not set. +func GetConfiguredDefaultProfileFrom(configFile *config.File) string { + section, err := configFile.GetSection(databricksSettingsSection) + if err != nil { + return "" + } + key, err := section.GetKey("default_profile") + if err != nil { + return "" + } + return key.String() +} + +// GetDefaultProfile returns the name of the default profile by loading the +// config file at configFilePath. Returns "" if the file doesn't exist. +// See GetDefaultProfileFrom for resolution order. +func GetDefaultProfile(ctx context.Context, configFilePath string) (string, error) { + configFile, err := loadConfigFile(ctx, configFilePath) + if err != nil { + return "", err + } + if configFile == nil { + return "", nil + } + return GetDefaultProfileFrom(configFile), nil +} + +// loadConfigFile loads a config file without creating it if it doesn't exist. +// Returns (nil, nil) when the file is not found. +func loadConfigFile(ctx context.Context, filename string) (*config.File, error) { + filename, err := resolveConfigFilePath(ctx, filename) + if err != nil { + return nil, err + } + configFile, err := config.LoadFile(filename) + if errors.Is(err, fs.ErrNotExist) { + return nil, nil + } + if err != nil { + return nil, fmt.Errorf("parse %s: %w", filename, err) + } + return configFile, nil +} + +// resolveConfigFilePath defaults to ~/.databrickscfg and expands ~ to the home directory. +func resolveConfigFilePath(ctx context.Context, filename string) (string, error) { if filename == "" { filename = "~/.databrickscfg" } - // Expand ~ to home directory, as we need a deterministic name for os.OpenFile - // to work in the cases when ~/.databrickscfg does not exist yet if strings.HasPrefix(filename, "~") { homedir, err := env.UserHomeDir(ctx) if err != nil { - return nil, fmt.Errorf("cannot find homedir: %w", err) + return "", fmt.Errorf("cannot find homedir: %w", err) } filename = fmt.Sprintf("%s%s", homedir, filename[1:]) } + return filename, nil +} + +// GetDefaultProfileFrom returns the name of the default profile from an +// already-loaded config file. It uses the following resolution order: +// 1. Explicit default_profile key in [__databricks-settings__]. +// 2. If there is exactly one profile in the file, return it. +// 3. If a profile named DEFAULT exists, return it. +// 4. Empty string (no default). +func GetDefaultProfileFrom(configFile *config.File) string { + // 1. Check for explicit default_profile setting. + if profile := GetConfiguredDefaultProfileFrom(configFile); profile != "" { + return profile + } + + // Collect profile sections (sections that have a "host" key, excluding + // the settings section). + var profileNames []string + hasDefault := false + for _, s := range configFile.Sections() { + if s.Name() == databricksSettingsSection { + continue + } + if !s.HasKey("host") { + continue + } + profileNames = append(profileNames, s.Name()) + if s.Name() == ini.DefaultSection { + hasDefault = true + } + } + + // 2. Exactly one profile: treat it as the default. + if len(profileNames) == 1 { + return profileNames[0] + } + + // 3. Legacy fallback: a DEFAULT section with a host key. + if hasDefault { + return ini.DefaultSection + } + + return "" +} + +// SetDefaultProfile writes the default_profile key to the [__databricks-settings__] section. +func SetDefaultProfile(ctx context.Context, profileName, configFilePath string) error { + configFile, err := loadOrCreateConfigFile(ctx, configFilePath) + if err != nil { + return err + } + + section, err := configFile.GetSection(databricksSettingsSection) + if err != nil { + // Section doesn't exist, create it. + section, err = configFile.NewSection(databricksSettingsSection) + if err != nil { + return fmt.Errorf("cannot create %s section: %w", databricksSettingsSection, err) + } + } + + section.Key("default_profile").SetValue(profileName) + + return backupAndSaveConfigFile(ctx, configFile) +} + +// backupAndSaveConfigFile adds a default section comment if needed, creates +// a .bak backup of the existing file, and saves the config file to disk. +func backupAndSaveConfigFile(ctx context.Context, configFile *config.File) error { + // Add a comment to the default section if it's empty. + section := configFile.Section(ini.DefaultSection) + if len(section.Keys()) == 0 && section.Comment == "" { + section.Comment = defaultComment + } + + orig, backupErr := os.ReadFile(configFile.Path()) + if len(orig) > 0 && backupErr == nil { + log.Infof(ctx, "Backing up in %s.bak", configFile.Path()) + err := os.WriteFile(configFile.Path()+".bak", orig, fileMode) + if err != nil { + return fmt.Errorf("backup: %w", err) + } + log.Infof(ctx, "Overwriting %s", configFile.Path()) + } else if backupErr != nil { + log.Warnf(ctx, "Failed to backup %s: %v. Proceeding to save", + configFile.Path(), backupErr) + } else { + log.Infof(ctx, "Saving %s", configFile.Path()) + } + return configFile.SaveTo(configFile.Path()) +} + +func loadOrCreateConfigFile(ctx context.Context, filename string) (*config.File, error) { + filename, err := resolveConfigFilePath(ctx, filename) + if err != nil { + return nil, err + } configFile, err := config.LoadFile(filename) if err != nil && errors.Is(err, fs.ErrNotExist) { file, err := os.OpenFile(filename, os.O_RDWR|os.O_CREATE|os.O_TRUNC, fileMode) @@ -130,27 +287,7 @@ func SaveToProfile(ctx context.Context, cfg *config.Config, clearKeys ...string) key.SetValue(attr.GetString(cfg)) } - // Add a comment to the default section if it's empty. - section = configFile.Section(ini.DefaultSection) - if len(section.Keys()) == 0 && section.Comment == "" { - section.Comment = defaultComment - } - - orig, backupErr := os.ReadFile(configFile.Path()) - if len(orig) > 0 && backupErr == nil { - log.Infof(ctx, "Backing up in %s.bak", configFile.Path()) - err = os.WriteFile(configFile.Path()+".bak", orig, fileMode) - if err != nil { - return fmt.Errorf("backup: %w", err) - } - log.Infof(ctx, "Overwriting %s", configFile.Path()) - } else if backupErr != nil { - log.Warnf(ctx, "Failed to backup %s: %v. Proceeding to save", - configFile.Path(), backupErr) - } else { - log.Infof(ctx, "Saving %s", configFile.Path()) - } - return configFile.SaveTo(configFile.Path()) + return backupAndSaveConfigFile(ctx, configFile) } func ValidateConfigAndProfileHost(cfg *config.Config, profile string) error { diff --git a/libs/databrickscfg/ops_test.go b/libs/databrickscfg/ops_test.go index 9032763bb9..9153c7ff2c 100644 --- a/libs/databrickscfg/ops_test.go +++ b/libs/databrickscfg/ops_test.go @@ -177,6 +177,201 @@ token = xyz `, string(contents)) } +func TestGetDefaultProfile(t *testing.T) { + testCases := []struct { + name string + content string + want string + }{ + { + name: "explicit default_profile setting", + content: "[__databricks-settings__]\ndefault_profile = my-workspace\n\n[my-workspace]\nhost = https://abc\n", + want: "my-workspace", + }, + { + name: "single profile fallback", + content: "[profile1]\nhost = https://abc\n", + want: "profile1", + }, + { + name: "multiple profiles no default", + content: "[profile1]\nhost = https://abc\n\n[profile2]\nhost = https://def\n", + want: "", + }, + { + name: "multiple profiles with DEFAULT fallback", + content: "[DEFAULT]\nhost = https://abc\n\n[profile2]\nhost = https://def\n", + want: "DEFAULT", + }, + { + name: "settings section without key single profile", + content: "[__databricks-settings__]\n\n[profile1]\nhost = https://abc\n", + want: "profile1", + }, + { + name: "empty config file", + content: "", + want: "", + }, + { + name: "settings section is not counted as a profile", + content: "[__databricks-settings__]\nsome_key = value\n\n[profile1]\nhost = https://abc\n", + want: "profile1", + }, + { + name: "section without host is not a profile", + content: "[no-host]\naccount_id = abc\n\n[profile1]\nhost = https://abc\n", + want: "profile1", + }, + } + + for _, tc := range testCases { + t.Run(tc.name, func(t *testing.T) { + path := filepath.Join(t.TempDir(), "databrickscfg") + err := os.WriteFile(path, []byte(tc.content), 0o600) + require.NoError(t, err) + + got, err := GetDefaultProfile(t.Context(), path) + require.NoError(t, err) + assert.Equal(t, tc.want, got) + }) + } +} + +func TestGetDefaultProfile_NoFile(t *testing.T) { + path := filepath.Join(t.TempDir(), "databrickscfg") + got, err := GetDefaultProfile(t.Context(), path) + require.NoError(t, err) + assert.Equal(t, "", got) + // Verify the file was NOT created as a side effect. + assert.NoFileExists(t, path) +} + +func TestGetConfiguredDefaultProfile(t *testing.T) { + testCases := []struct { + name string + content string + want string + }{ + { + name: "explicit default_profile setting", + content: "[__databricks-settings__]\ndefault_profile = my-workspace\n\n[my-workspace]\nhost = https://abc\n", + want: "my-workspace", + }, + { + name: "single profile fallback is ignored", + content: "[profile1]\nhost = https://abc\n", + want: "", + }, + { + name: "DEFAULT fallback is ignored", + content: "[DEFAULT]\nhost = https://abc\n\n[profile2]\nhost = https://def\n", + want: "", + }, + { + name: "settings section without key", + content: "[__databricks-settings__]\n\n[profile1]\nhost = https://abc\n", + want: "", + }, + } + + for _, tc := range testCases { + t.Run(tc.name, func(t *testing.T) { + path := filepath.Join(t.TempDir(), "databrickscfg") + err := os.WriteFile(path, []byte(tc.content), 0o600) + require.NoError(t, err) + + got, err := GetConfiguredDefaultProfile(t.Context(), path) + require.NoError(t, err) + assert.Equal(t, tc.want, got) + }) + } +} + +func TestGetConfiguredDefaultProfile_NoFile(t *testing.T) { + path := filepath.Join(t.TempDir(), "databrickscfg") + got, err := GetConfiguredDefaultProfile(t.Context(), path) + require.NoError(t, err) + assert.Equal(t, "", got) + // Verify the file was NOT created as a side effect. + assert.NoFileExists(t, path) +} + +func TestSetDefaultProfile(t *testing.T) { + testCases := []struct { + name string + initial string + profile string + wantKey string + }{ + { + name: "creates section and key", + initial: "[profile1]\nhost = https://abc\n", + profile: "profile1", + wantKey: "profile1", + }, + { + name: "updates existing key", + initial: "[__databricks-settings__]\ndefault_profile = old-profile\n\n[profile1]\nhost = https://abc\n", + profile: "new-profile", + wantKey: "new-profile", + }, + { + name: "creates section in empty file", + initial: "", + profile: "my-workspace", + wantKey: "my-workspace", + }, + } + + for _, tc := range testCases { + t.Run(tc.name, func(t *testing.T) { + ctx := t.Context() + path := filepath.Join(t.TempDir(), "databrickscfg") + err := os.WriteFile(path, []byte(tc.initial), 0o600) + require.NoError(t, err) + + err = SetDefaultProfile(ctx, tc.profile, path) + require.NoError(t, err) + + got, err := GetDefaultProfile(ctx, path) + require.NoError(t, err) + assert.Equal(t, tc.wantKey, got) + }) + } +} + +func TestSetDefaultProfile_RoundTrip(t *testing.T) { + ctx := t.Context() + path := filepath.Join(t.TempDir(), "databrickscfg") + + // Start with a profile. + err := SaveToProfile(ctx, &config.Config{ + ConfigFile: path, + Profile: "my-workspace", + Host: "https://abc.cloud.databricks.com", + Token: "xyz", + }) + require.NoError(t, err) + + // Set it as default. + err = SetDefaultProfile(ctx, "my-workspace", path) + require.NoError(t, err) + + // Read it back. + got, err := GetDefaultProfile(ctx, path) + require.NoError(t, err) + assert.Equal(t, "my-workspace", got) + + // Verify the profile section is still intact. + file, err := loadOrCreateConfigFile(ctx, path) + require.NoError(t, err) + section, err := file.GetSection("my-workspace") + require.NoError(t, err) + assert.Equal(t, "https://abc.cloud.databricks.com", section.Key("host").String()) + assert.Equal(t, "xyz", section.Key("token").String()) +} + func TestSaveToProfile_MergeSemantics(t *testing.T) { type saveOp struct { cfg *config.Config