diff --git a/acceptance/cmd/auth/login/preserve-fields/out.databrickscfg b/acceptance/cmd/auth/login/preserve-fields/out.databrickscfg new file mode 100644 index 0000000000..3658c4d33b --- /dev/null +++ b/acceptance/cmd/auth/login/preserve-fields/out.databrickscfg @@ -0,0 +1,7 @@ +[DEFAULT] +host = [DATABRICKS_URL] +cluster_id = existing-cluster-123 +warehouse_id = warehouse-456 +azure_environment = USGOVERNMENT +custom_key = my-custom-value +auth_type = databricks-cli diff --git a/acceptance/cmd/auth/login/preserve-fields/out.test.toml b/acceptance/cmd/auth/login/preserve-fields/out.test.toml new file mode 100644 index 0000000000..d560f1de04 --- /dev/null +++ b/acceptance/cmd/auth/login/preserve-fields/out.test.toml @@ -0,0 +1,5 @@ +Local = true +Cloud = false + +[EnvMatrix] + DATABRICKS_BUNDLE_ENGINE = ["terraform", "direct"] diff --git a/acceptance/cmd/auth/login/preserve-fields/output.txt b/acceptance/cmd/auth/login/preserve-fields/output.txt new file mode 100644 index 0000000000..625efd0bf2 --- /dev/null +++ b/acceptance/cmd/auth/login/preserve-fields/output.txt @@ -0,0 +1,21 @@ + +=== Initial profile +[DEFAULT] +host = [DATABRICKS_URL] +cluster_id = existing-cluster-123 +warehouse_id = warehouse-456 +azure_environment = USGOVERNMENT +custom_key = my-custom-value + +=== Run auth login (no --configure-cluster or --configure-serverless) +>>> [CLI] auth login --host [DATABRICKS_URL] --profile DEFAULT +Profile DEFAULT was successfully saved + +=== Profile after auth login — all non-auth fields should be preserved +[DEFAULT] +host = [DATABRICKS_URL] +cluster_id = existing-cluster-123 +warehouse_id = warehouse-456 +azure_environment = USGOVERNMENT +custom_key = my-custom-value +auth_type = databricks-cli diff --git a/acceptance/cmd/auth/login/preserve-fields/script b/acceptance/cmd/auth/login/preserve-fields/script new file mode 100644 index 0000000000..a82101c892 --- /dev/null +++ b/acceptance/cmd/auth/login/preserve-fields/script @@ -0,0 +1,28 @@ +sethome "./home" + +# Create an initial profile with cluster_id, warehouse_id, azure_environment, +# and a custom key that is not a recognized SDK config attribute. +cat > "./home/.databrickscfg" <>> [CLI] configure --token --host https://host + +=== Profile after PAT configure +[DEFAULT] +host = https://host +account_id = acc-123 +workspace_id = ws-456 +token = [DATABRICKS_TOKEN] diff --git a/acceptance/cmd/configure/clears-oauth-on-pat/script b/acceptance/cmd/configure/clears-oauth-on-pat/script new file mode 100644 index 0000000000..587e0d0a55 --- /dev/null +++ b/acceptance/cmd/configure/clears-oauth-on-pat/script @@ -0,0 +1,24 @@ +sethome "./home" + +# Pre-populate a profile with OAuth metadata and unified host fields +# (as if `auth login` was previously run against a unified host). +cat > "./home/.databrickscfg" <>> [CLI] configure --token --host https://host + +=== Profile after configure (serverless should be cleared) +[DEFAULT] +host = https://host +cluster_id = env-cluster-789 +token = [DATABRICKS_TOKEN] diff --git a/acceptance/cmd/configure/clears-serverless-when-cluster-from-env/script b/acceptance/cmd/configure/clears-serverless-when-cluster-from-env/script new file mode 100644 index 0000000000..b9d8573f78 --- /dev/null +++ b/acceptance/cmd/configure/clears-serverless-when-cluster-from-env/script @@ -0,0 +1,22 @@ +sethome "./home" + +# Pre-populate a profile with serverless_compute_id. +cat > "./home/.databrickscfg" < 0 { + persistentAuthOpts = append(persistentAuthOpts, u2m.WithScopes(scopesList)) + } + persistentAuth, err := u2m.NewPersistentAuth(ctx, persistentAuthOpts...) if err != nil { return "", nil, err } @@ -443,6 +456,10 @@ func runInlineLogin(ctx context.Context, profiler profile.Profiler) (string, *pr return "", nil, err } + clearKeys := oauthLoginClearKeys() + if !loginArgs.IsUnifiedHost { + clearKeys = append(clearKeys, "experimental_is_unified_host") + } err = databrickscfg.SaveToProfile(ctx, &config.Config{ Profile: profileName, Host: loginArgs.Host, @@ -451,7 +468,8 @@ func runInlineLogin(ctx context.Context, profiler profile.Profiler) (string, *pr WorkspaceID: loginArgs.WorkspaceID, Experimental_IsUnifiedHost: loginArgs.IsUnifiedHost, ConfigFile: os.Getenv("DATABRICKS_CONFIG_FILE"), - }) + Scopes: scopesList, + }, clearKeys...) if err != nil { return "", nil, err } diff --git a/cmd/configure/configure.go b/cmd/configure/configure.go index 4f36bf7c72..f134c91191 100644 --- a/cmd/configure/configure.go +++ b/cmd/configure/configure.go @@ -12,6 +12,15 @@ import ( "github.com/spf13/cobra" ) +// patConfigureExtraClearKeys lists non-credential profile keys that should also +// be cleared when saving a PAT-based profile. Auth credential keys are derived +// dynamically from config.ConfigAttributes via databrickscfg.AuthCredentialKeys(). +var patConfigureExtraClearKeys = []string{ + "auth_type", + "scopes", + "databricks_cli_path", +} + func configureInteractive(cmd *cobra.Command, flags *configureFlags, cfg *config.Config) error { ctx := cmd.Context() @@ -141,14 +150,28 @@ The host must be specified with the --host flag or the DATABRICKS_HOST environme // This is relevant for OAuth only. cfg.DatabricksCliPath = "" - // Save profile to config file. + // Save profile to config file. PAT-based configure clears all + // non-PAT auth credentials and OAuth metadata to prevent + // multi-auth conflicts in the profile. + clearKeys := append(databrickscfg.AuthCredentialKeys(), patConfigureExtraClearKeys...) + + // Cluster and serverless are mutually exclusive. Clear serverless + // when a cluster is being set (via flag or env var). + if cfg.ClusterID != "" { + clearKeys = append(clearKeys, "serverless_compute_id") + } + + // Clear stale unified-host metadata — PAT profiles don't use it, + // and leaving it can change HostType() routing. + clearKeys = append(clearKeys, "experimental_is_unified_host") + return databrickscfg.SaveToProfile(ctx, &config.Config{ Profile: cfg.Profile, Host: cfg.Host, Token: cfg.Token, ClusterID: cfg.ClusterID, ConfigFile: cfg.ConfigFile, - }) + }, clearKeys...) } return cmd diff --git a/libs/databrickscfg/ops.go b/libs/databrickscfg/ops.go index 32c53b0004..e38a0bd69b 100644 --- a/libs/databrickscfg/ops.go +++ b/libs/databrickscfg/ops.go @@ -80,7 +80,26 @@ func matchOrCreateSection(ctx context.Context, configFile *config.File, cfg *con return section, nil } -func SaveToProfile(ctx context.Context, cfg *config.Config) error { +// AuthCredentialKeys returns the config file key names for all auth credential +// fields from the SDK's ConfigAttributes. These are fields annotated with an +// auth type (e.g. pat, basic, oauth, azure, google). Use this to clear stale +// credentials when switching auth methods. +func AuthCredentialKeys() []string { + var keys []string + for _, attr := range config.ConfigAttributes { + if attr.HasAuthAttribute() { + keys = append(keys, attr.Name) + } + } + return keys +} + +// SaveToProfile merges the provided config into a .databrickscfg profile. +// Non-zero fields in cfg overwrite existing values. Existing keys not +// mentioned in cfg are preserved. Keys listed in clearKeys are explicitly +// removed (use this for mutually exclusive fields like cluster_id vs +// serverless_compute_id, or to drop stale auth credentials on auth-type switch). +func SaveToProfile(ctx context.Context, cfg *config.Config, clearKeys ...string) error { configFile, err := loadOrCreateConfigFile(cfg.ConfigFile) if err != nil { return err @@ -95,11 +114,13 @@ func SaveToProfile(ctx context.Context, cfg *config.Config) error { cfg.Profile = "" cfg.ConfigFile = "" - // clear old keys in case we're overriding the section - for _, oldKey := range section.KeyStrings() { - section.DeleteKey(oldKey) + // Explicitly remove keys the caller wants cleared. + for _, key := range clearKeys { + section.DeleteKey(key) } + // Write non-zero fields from the new config. Iterates ConfigAttributes + // in declaration order for deterministic key ordering on new profiles. for _, attr := range config.ConfigAttributes { if attr.IsZero(cfg) { continue diff --git a/libs/databrickscfg/ops_test.go b/libs/databrickscfg/ops_test.go index ed81da2d10..646ed1e54a 100644 --- a/libs/databrickscfg/ops_test.go +++ b/libs/databrickscfg/ops_test.go @@ -178,74 +178,103 @@ token = xyz `, string(contents)) } -func TestSaveToProfile_ClearingPreviousProfile(t *testing.T) { - ctx := context.Background() - path := filepath.Join(t.TempDir(), "databrickscfg") - - err := SaveToProfile(ctx, &config.Config{ - ConfigFile: path, - Profile: "abc", - Host: "https://foo", - Token: "xyz", - }) - assert.NoError(t, err) - - err = SaveToProfile(ctx, &config.Config{ - ConfigFile: path, - Profile: "bcd", - Host: "https://bar", - Token: "zyx", - }) - assert.NoError(t, err) - assert.FileExists(t, path+".bak") - - err = SaveToProfile(ctx, &config.Config{ - ConfigFile: path, - Host: "https://foo", - AuthType: "databricks-cli", - }) - assert.NoError(t, err) - - file, err := loadOrCreateConfigFile(path) - assert.NoError(t, err) - - assert.Len(t, file.Sections(), 3) - assert.True(t, file.HasSection("DEFAULT")) - assert.True(t, file.HasSection("bcd")) - assert.True(t, file.HasSection("bcd")) - - dlft, err := file.GetSection("DEFAULT") - assert.NoError(t, err) - assert.Empty(t, dlft.KeysHash()) - - abc, err := file.GetSection("abc") - assert.NoError(t, err) - raw := abc.KeysHash() - assert.Len(t, raw, 2) - assert.Equal(t, "https://foo", raw["host"]) - assert.Equal(t, "databricks-cli", raw["auth_type"]) -} - -func TestSaveToProfile_WithScopes(t *testing.T) { - ctx := context.Background() - path := filepath.Join(t.TempDir(), "databrickscfg") +func TestSaveToProfile_MergeSemantics(t *testing.T) { + type saveOp struct { + cfg *config.Config + clearKeys []string + } - err := SaveToProfile(ctx, &config.Config{ - ConfigFile: path, - Profile: "scoped", - Host: "https://myworkspace.cloud.databricks.com", - AuthType: "databricks-cli", - Scopes: []string{"jobs", "pipelines", "clusters"}, - }) - require.NoError(t, err) + testCases := []struct { + name string + profile string + saves []saveOp + wantKeys map[string]string + }{ + { + name: "preserves existing keys on merge", + profile: "abc", + saves: []saveOp{ + {cfg: &config.Config{Profile: "abc", Host: "https://foo", Token: "xyz"}}, + {cfg: &config.Config{Host: "https://foo", AuthType: "databricks-cli"}}, + }, + wantKeys: map[string]string{ + "host": "https://foo", + "auth_type": "databricks-cli", + "token": "xyz", + }, + }, + { + name: "clear keys removes specified keys", + profile: "abc", + saves: []saveOp{ + {cfg: &config.Config{Profile: "abc", Host: "https://foo", Token: "xyz", ClusterID: "cluster-123"}}, + {cfg: &config.Config{Host: "https://foo", AuthType: "databricks-cli", ServerlessComputeID: "auto"}, clearKeys: []string{"token", "cluster_id"}}, + }, + wantKeys: map[string]string{ + "host": "https://foo", + "auth_type": "databricks-cli", + "serverless_compute_id": "auto", + }, + }, + { + name: "overwrites existing values", + profile: "abc", + saves: []saveOp{ + {cfg: &config.Config{Profile: "abc", Host: "https://old-host"}}, + {cfg: &config.Config{Profile: "abc", Host: "https://new-host"}}, + }, + wantKeys: map[string]string{ + "host": "https://new-host", + }, + }, + { + name: "clear nonexistent key is noop", + profile: "abc", + saves: []saveOp{ + {cfg: &config.Config{Profile: "abc", Host: "https://foo"}}, + {cfg: &config.Config{Profile: "abc", Host: "https://foo", AuthType: "databricks-cli"}, clearKeys: []string{"token", "nonexistent_key"}}, + }, + wantKeys: map[string]string{ + "host": "https://foo", + "auth_type": "databricks-cli", + }, + }, + { + name: "writes scopes as comma-separated", + profile: "scoped", + saves: []saveOp{ + {cfg: &config.Config{Profile: "scoped", Host: "https://myworkspace.cloud.databricks.com", AuthType: "databricks-cli", Scopes: []string{"jobs", "pipelines", "clusters"}}}, + }, + wantKeys: map[string]string{ + "host": "https://myworkspace.cloud.databricks.com", + "auth_type": "databricks-cli", + "scopes": "jobs,pipelines,clusters", + }, + }, + } - file, err := loadOrCreateConfigFile(path) - require.NoError(t, err) - section, err := file.GetSection("scoped") - require.NoError(t, err) - raw := section.KeysHash() - assert.Len(t, raw, 3) - assert.Equal(t, "https://myworkspace.cloud.databricks.com", raw["host"]) - assert.Equal(t, "databricks-cli", raw["auth_type"]) - assert.Equal(t, "jobs,pipelines,clusters", raw["scopes"]) + for _, tc := range testCases { + t.Run(tc.name, func(t *testing.T) { + ctx := context.Background() + path := filepath.Join(t.TempDir(), "databrickscfg") + + for _, save := range tc.saves { + save.cfg.ConfigFile = path + err := SaveToProfile(ctx, save.cfg, save.clearKeys...) + require.NoError(t, err) + } + + file, err := loadOrCreateConfigFile(path) + require.NoError(t, err) + + section, err := file.GetSection(tc.profile) + require.NoError(t, err) + + raw := section.KeysHash() + assert.Len(t, raw, len(tc.wantKeys)) + for k, v := range tc.wantKeys { + assert.Equal(t, v, raw[k], "key %s", k) + } + }) + } } diff --git a/libs/databrickscfg/profile/file.go b/libs/databrickscfg/profile/file.go index b078e67d3a..3d062e11af 100644 --- a/libs/databrickscfg/profile/file.go +++ b/libs/databrickscfg/profile/file.go @@ -87,6 +87,7 @@ func (f FileProfilerImpl) LoadProfiles(ctx context.Context, fn ProfileMatchFunct ClusterID: all["cluster_id"], ServerlessComputeID: all["serverless_compute_id"], HasClientCredentials: all["client_id"] != "" && all["client_secret"] != "", + Scopes: all["scopes"], } if fn(profile) { profiles = append(profiles, profile) diff --git a/libs/databrickscfg/profile/profile.go b/libs/databrickscfg/profile/profile.go index 9cdf24f82b..7bf8784ad2 100644 --- a/libs/databrickscfg/profile/profile.go +++ b/libs/databrickscfg/profile/profile.go @@ -18,6 +18,7 @@ type Profile struct { ClusterID string ServerlessComputeID string HasClientCredentials bool + Scopes string } func (p Profile) Cloud() string {