From ce79a102f673ec83215e439d2da4d0e4d567541a Mon Sep 17 00:00:00 2001 From: simon Date: Fri, 13 Mar 2026 00:26:18 +0100 Subject: [PATCH 1/7] Add text output mode for auth token and auth env commands Both commands now respect --output text when explicitly set: - auth token --output text: outputs just the access token string - auth env --output text: outputs KEY=VALUE lines JSON remains the default for backward compatibility. Co-authored-by: Isaac --- cmd/auth/env.go | 64 +++++++++++++++++++--- cmd/auth/env_test.go | 30 +++++++++++ cmd/auth/token.go | 9 ++++ cmd/auth/token_test.go | 117 +++++++++++++++++++++++++++++++++++++++++ 4 files changed, 212 insertions(+), 8 deletions(-) create mode 100644 cmd/auth/env_test.go diff --git a/cmd/auth/env.go b/cmd/auth/env.go index 11149af8c0..9b75204e9d 100644 --- a/cmd/auth/env.go +++ b/cmd/auth/env.go @@ -10,7 +10,9 @@ import ( "net/url" "strings" + "github.com/databricks/cli/cmd/root" "github.com/databricks/cli/libs/databrickscfg/profile" + "github.com/databricks/cli/libs/flags" "github.com/databricks/databricks-sdk-go/config" "github.com/spf13/cobra" "gopkg.in/ini.v1" @@ -122,16 +124,23 @@ func newEnvCommand() *cobra.Command { if err != nil { return err } - vars := map[string]string{} - for _, a := range config.ConfigAttributes { - if a.IsZero(cfg) { - continue - } - envValue := a.GetString(cfg) - for _, envName := range a.EnvVars { - vars[envName] = envValue + vars := collectEnvVars(cfg) + + // Output KEY=VALUE lines when the user explicitly passes --output text. + if cmd.Flag("output").Changed && root.OutputType(cmd) == flags.OutputText { + w := cmd.OutOrStdout() + for _, a := range config.ConfigAttributes { + if a.IsZero(cfg) { + continue + } + v := a.GetString(cfg) + for _, envName := range a.EnvVars { + fmt.Fprintf(w, "%s=%s\n", envName, quoteEnvValue(v)) + } } + return nil } + raw, err := json.MarshalIndent(map[string]any{ "env": vars, }, "", " ") @@ -144,3 +153,42 @@ func newEnvCommand() *cobra.Command { return cmd } + +// collectEnvVars returns the environment variables for the given config +// as a map from env var name to value. +func collectEnvVars(cfg *config.Config) map[string]string { + vars := map[string]string{} + for _, a := range config.ConfigAttributes { + if a.IsZero(cfg) { + continue + } + v := a.GetString(cfg) + for _, envName := range a.EnvVars { + vars[envName] = v + } + } + return vars +} + +// quoteEnvValue quotes a value for KEY=VALUE output if it contains spaces, +// double quotes, or shell-special characters. Embedded double quotes and +// backslashes are escaped with a backslash. +func quoteEnvValue(v string) string { + if v == "" { + return `""` + } + needsQuoting := strings.ContainsAny(v, " \t\"\\$`!#&|;(){}[]<>?*~'") + if !needsQuoting { + return v + } + var b strings.Builder + b.WriteByte('"') + for _, c := range v { + if c == '"' || c == '\\' { + b.WriteByte('\\') + } + b.WriteRune(c) + } + b.WriteByte('"') + return b.String() +} diff --git a/cmd/auth/env_test.go b/cmd/auth/env_test.go new file mode 100644 index 0000000000..ac6ab93b6f --- /dev/null +++ b/cmd/auth/env_test.go @@ -0,0 +1,30 @@ +package auth + +import ( + "testing" + + "github.com/stretchr/testify/assert" +) + +func TestQuoteEnvValue(t *testing.T) { + cases := []struct { + name string + in string + want string + }{ + {name: "simple value", in: "hello", want: "hello"}, + {name: "empty value", in: "", want: `""`}, + {name: "value with space", in: "hello world", want: `"hello world"`}, + {name: "value with tab", in: "hello\tworld", want: "\"hello\tworld\""}, + {name: "value with double quote", in: `say "hi"`, want: `"say \"hi\""`}, + {name: "value with backslash", in: `path\to`, want: `"path\\to"`}, + {name: "url value", in: "https://example.com", want: "https://example.com"}, + {name: "value with dollar", in: "price$5", want: `"price$5"`}, + } + for _, c := range cases { + t.Run(c.name, func(t *testing.T) { + got := quoteEnvValue(c.in) + assert.Equal(t, c.want, got) + }) + } +} diff --git a/cmd/auth/token.go b/cmd/auth/token.go index b33722c1ed..687235ee2f 100644 --- a/cmd/auth/token.go +++ b/cmd/auth/token.go @@ -8,11 +8,13 @@ import ( "strings" "time" + "github.com/databricks/cli/cmd/root" "github.com/databricks/cli/libs/auth" "github.com/databricks/cli/libs/cmdio" "github.com/databricks/cli/libs/databrickscfg" "github.com/databricks/cli/libs/databrickscfg/profile" "github.com/databricks/cli/libs/env" + "github.com/databricks/cli/libs/flags" "github.com/databricks/databricks-sdk-go/config" "github.com/databricks/databricks-sdk-go/credentials/u2m" "github.com/databricks/databricks-sdk-go/credentials/u2m/cache" @@ -83,6 +85,13 @@ using a client ID and secret is not supported.`, if err != nil { return err } + + // Output plain token when the user explicitly passes --output text. + if cmd.Flag("output").Changed && root.OutputType(cmd) == flags.OutputText { + _, _ = fmt.Fprint(cmd.OutOrStdout(), t.AccessToken) + return nil + } + raw, err := json.MarshalIndent(t, "", " ") if err != nil { return err diff --git a/cmd/auth/token_test.go b/cmd/auth/token_test.go index aa343eb372..32d19fac69 100644 --- a/cmd/auth/token_test.go +++ b/cmd/auth/token_test.go @@ -1,17 +1,23 @@ package auth import ( + "bytes" "context" + "encoding/json" + "fmt" "net/http" "testing" "time" + "github.com/databricks/cli/cmd/root" "github.com/databricks/cli/libs/auth" "github.com/databricks/cli/libs/cmdio" "github.com/databricks/cli/libs/databrickscfg/profile" "github.com/databricks/cli/libs/env" + "github.com/databricks/cli/libs/flags" "github.com/databricks/databricks-sdk-go/credentials/u2m" "github.com/databricks/databricks-sdk-go/httpclient/fixtures" + "github.com/spf13/cobra" "github.com/stretchr/testify/assert" "golang.org/x/oauth2" ) @@ -729,3 +735,114 @@ func (e errProfiler) LoadProfiles(context.Context, profile.ProfileMatchFunction) func (e errProfiler) GetPath(context.Context) (string, error) { return "", nil } + +func TestTokenCommand_TextOutput(t *testing.T) { + profiler := profile.InMemoryProfiler{ + Profiles: profile.Profiles{ + { + Name: "test-ws", + Host: "https://test-ws.cloud.databricks.com", + }, + }, + } + tokenCache := &inMemoryTokenCache{ + Tokens: map[string]*oauth2.Token{ + "test-ws": { + RefreshToken: "test-ws", + Expiry: time.Now().Add(1 * time.Hour), + }, + }, + } + persistentAuthOpts := []u2m.PersistentAuthOption{ + u2m.WithTokenCache(tokenCache), + u2m.WithOAuthEndpointSupplier(&MockApiClient{}), + u2m.WithHttpClient(&http.Client{Transport: fixtures.SliceTransport{refreshSuccessTokenResponse}}), + } + + cases := []struct { + name string + args []string + wantSubstr string + wantJSON bool + }{ + { + name: "default output is JSON", + args: []string{"--profile", "test-ws"}, + wantSubstr: `"access_token"`, + wantJSON: true, + }, + { + name: "explicit --output json produces JSON", + args: []string{"--profile", "test-ws", "--output", "json"}, + wantSubstr: `"access_token"`, + wantJSON: true, + }, + { + name: "explicit --output text produces plain token", + args: []string{"--profile", "test-ws", "--output", "text"}, + wantSubstr: "new-access-token", + wantJSON: false, + }, + } + + for _, c := range cases { + t.Run(c.name, func(t *testing.T) { + ctx := cmdio.MockDiscard(t.Context()) + authArgs := &auth.AuthArguments{} + + parent := &cobra.Command{Use: "databricks"} + outputFlag := flags.OutputText + parent.PersistentFlags().VarP(&outputFlag, "output", "o", "output type: text or json") + parent.PersistentFlags().StringP("profile", "p", "", "~/.databrickscfg profile") + + tokenCmd := newTokenCommand(authArgs) + // Override RunE to inject test profiler and token cache while + // keeping the same output formatting logic as the real command. + tokenCmd.RunE = func(cmd *cobra.Command, args []string) error { + profileName := "" + if f := cmd.Flag("profile"); f != nil { + profileName = f.Value.String() + } + tok, err := loadToken(cmd.Context(), loadTokenArgs{ + authArguments: authArgs, + profileName: profileName, + args: args, + tokenTimeout: 1 * time.Hour, + profiler: profiler, + persistentAuthOpts: persistentAuthOpts, + }) + if err != nil { + return err + } + if cmd.Flag("output").Changed && root.OutputType(cmd) == flags.OutputText { + _, _ = fmt.Fprint(cmd.OutOrStdout(), tok.AccessToken) + return nil + } + raw, err := json.MarshalIndent(tok, "", " ") + if err != nil { + return err + } + _, _ = cmd.OutOrStdout().Write(raw) + return nil + } + + parent.AddCommand(tokenCmd) + parent.SetContext(ctx) + + var buf bytes.Buffer + parent.SetOut(&buf) + parent.SetArgs(append([]string{"token"}, c.args...)) + + err := parent.Execute() + assert.NoError(t, err) + + output := buf.String() + assert.Contains(t, output, c.wantSubstr) + if c.wantJSON { + assert.Contains(t, output, "{") + } else { + assert.NotContains(t, output, "{") + } + }) + } +} From aeae5b699cc9c6e909cc8f86c562f6aeead68365 Mon Sep 17 00:00:00 2001 From: simon Date: Fri, 13 Mar 2026 03:57:48 +0100 Subject: [PATCH 2/7] Fix review findings: shell escaping, test coverage, nits Co-authored-by: Isaac --- cmd/auth/env.go | 5 ++-- cmd/auth/env_test.go | 61 +++++++++++++++++++++++++++++++++++++++++- cmd/auth/token.go | 2 +- cmd/auth/token_test.go | 11 ++++---- 4 files changed, 69 insertions(+), 10 deletions(-) diff --git a/cmd/auth/env.go b/cmd/auth/env.go index 9b75204e9d..b73b21f577 100644 --- a/cmd/auth/env.go +++ b/cmd/auth/env.go @@ -124,8 +124,6 @@ func newEnvCommand() *cobra.Command { if err != nil { return err } - vars := collectEnvVars(cfg) - // Output KEY=VALUE lines when the user explicitly passes --output text. if cmd.Flag("output").Changed && root.OutputType(cmd) == flags.OutputText { w := cmd.OutOrStdout() @@ -141,6 +139,7 @@ func newEnvCommand() *cobra.Command { return nil } + vars := collectEnvVars(cfg) raw, err := json.MarshalIndent(map[string]any{ "env": vars, }, "", " ") @@ -184,7 +183,7 @@ func quoteEnvValue(v string) string { var b strings.Builder b.WriteByte('"') for _, c := range v { - if c == '"' || c == '\\' { + if c == '"' || c == '\\' || c == '$' || c == '`' || c == '!' { b.WriteByte('\\') } b.WriteRune(c) diff --git a/cmd/auth/env_test.go b/cmd/auth/env_test.go index ac6ab93b6f..c094a33e25 100644 --- a/cmd/auth/env_test.go +++ b/cmd/auth/env_test.go @@ -1,9 +1,13 @@ package auth import ( + "bytes" "testing" + "github.com/databricks/cli/libs/flags" + "github.com/spf13/cobra" "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" ) func TestQuoteEnvValue(t *testing.T) { @@ -19,7 +23,9 @@ func TestQuoteEnvValue(t *testing.T) { {name: "value with double quote", in: `say "hi"`, want: `"say \"hi\""`}, {name: "value with backslash", in: `path\to`, want: `"path\\to"`}, {name: "url value", in: "https://example.com", want: "https://example.com"}, - {name: "value with dollar", in: "price$5", want: `"price$5"`}, + {name: "value with dollar", in: "price$5", want: `"price\$5"`}, + {name: "value with backtick", in: "hello`world", want: `"hello\` + "`" + `world"`}, + {name: "value with bang", in: "hello!world", want: `"hello\!world"`}, } for _, c := range cases { t.Run(c.name, func(t *testing.T) { @@ -28,3 +34,56 @@ func TestQuoteEnvValue(t *testing.T) { }) } } + +func TestEnvCommand_TextOutput(t *testing.T) { + cases := []struct { + name string + args []string + wantJSON bool + }{ + { + name: "default output is JSON", + args: []string{"--host", "https://test.cloud.databricks.com"}, + wantJSON: true, + }, + { + name: "explicit --output text produces KEY=VALUE lines", + args: []string{"--host", "https://test.cloud.databricks.com", "--output", "text"}, + wantJSON: false, + }, + } + + for _, c := range cases { + t.Run(c.name, func(t *testing.T) { + parent := &cobra.Command{Use: "databricks"} + outputFlag := flags.OutputText + parent.PersistentFlags().VarP(&outputFlag, "output", "o", "output type: text or json") + + envCmd := newEnvCommand() + parent.AddCommand(envCmd) + + // Set DATABRICKS_TOKEN so the SDK's config.Authenticate succeeds + // without hitting a real endpoint. + t.Setenv("DATABRICKS_TOKEN", "test-token-value") + + var buf bytes.Buffer + parent.SetOut(&buf) + parent.SetArgs(append([]string{"env"}, c.args...)) + + err := parent.Execute() + require.NoError(t, err) + + output := buf.String() + if c.wantJSON { + assert.Contains(t, output, "{") + assert.Contains(t, output, "DATABRICKS_HOST") + } else { + assert.NotContains(t, output, "{") + assert.Contains(t, output, "DATABRICKS_HOST=") + assert.Contains(t, output, "=") + // Verify KEY=VALUE format (no JSON structure) + assert.NotContains(t, output, `"env"`) + } + }) + } +} diff --git a/cmd/auth/token.go b/cmd/auth/token.go index 687235ee2f..c8fb626f7d 100644 --- a/cmd/auth/token.go +++ b/cmd/auth/token.go @@ -88,7 +88,7 @@ using a client ID and secret is not supported.`, // Output plain token when the user explicitly passes --output text. if cmd.Flag("output").Changed && root.OutputType(cmd) == flags.OutputText { - _, _ = fmt.Fprint(cmd.OutOrStdout(), t.AccessToken) + _, _ = fmt.Fprintln(cmd.OutOrStdout(), t.AccessToken) return nil } diff --git a/cmd/auth/token_test.go b/cmd/auth/token_test.go index 32d19fac69..289313bae8 100644 --- a/cmd/auth/token_test.go +++ b/cmd/auth/token_test.go @@ -778,9 +778,9 @@ func TestTokenCommand_TextOutput(t *testing.T) { wantJSON: true, }, { - name: "explicit --output text produces plain token", + name: "explicit --output text produces plain token with newline", args: []string{"--profile", "test-ws", "--output", "text"}, - wantSubstr: "new-access-token", + wantSubstr: "new-access-token\n", wantJSON: false, }, } @@ -796,8 +796,9 @@ func TestTokenCommand_TextOutput(t *testing.T) { parent.PersistentFlags().StringP("profile", "p", "", "~/.databrickscfg profile") tokenCmd := newTokenCommand(authArgs) - // Override RunE to inject test profiler and token cache while - // keeping the same output formatting logic as the real command. + // Override RunE to inject test profiler and token cache. + // The output formatting logic below must mirror newTokenCommand.RunE. + // If you change the output logic in newTokenCommand, update this too. tokenCmd.RunE = func(cmd *cobra.Command, args []string) error { profileName := "" if f := cmd.Flag("profile"); f != nil { @@ -815,7 +816,7 @@ func TestTokenCommand_TextOutput(t *testing.T) { return err } if cmd.Flag("output").Changed && root.OutputType(cmd) == flags.OutputText { - _, _ = fmt.Fprint(cmd.OutOrStdout(), tok.AccessToken) + _, _ = fmt.Fprintln(cmd.OutOrStdout(), tok.AccessToken) return nil } raw, err := json.MarshalIndent(tok, "", " ") From 0a4a1c9ea36bae8a12c268a187e4ad69623c54c8 Mon Sep 17 00:00:00 2001 From: simon Date: Fri, 13 Mar 2026 07:24:47 +0100 Subject: [PATCH 3/7] Switch to single-quote shell escaping and improve test isolation --- cmd/auth/env.go | 23 ++++++++--------------- cmd/auth/env_test.go | 24 ++++++++++++++++-------- cmd/auth/token.go | 25 ++++++++++++++----------- cmd/auth/token_test.go | 19 +++---------------- 4 files changed, 41 insertions(+), 50 deletions(-) diff --git a/cmd/auth/env.go b/cmd/auth/env.go index b73b21f577..d0ef95c474 100644 --- a/cmd/auth/env.go +++ b/cmd/auth/env.go @@ -32,6 +32,8 @@ func canonicalHost(host string) (string, error) { var ErrNoMatchingProfiles = errors.New("no matching profiles found") +const shellQuotedSpecialChars = " \t\"\\$`!#&|;(){}[]<>?*~'" + func resolveSection(cfg *config.Config, iniFile *config.File) (*ini.Section, error) { var candidates []*ini.Section configuredHost, err := canonicalHost(cfg.Host) @@ -169,25 +171,16 @@ func collectEnvVars(cfg *config.Config) map[string]string { return vars } -// quoteEnvValue quotes a value for KEY=VALUE output if it contains spaces, -// double quotes, or shell-special characters. Embedded double quotes and -// backslashes are escaped with a backslash. +// quoteEnvValue quotes a value for KEY=VALUE output if it contains spaces or +// shell-special characters. Single quotes prevent shell expansion, and +// embedded single quotes use the POSIX-compatible '\'' sequence. func quoteEnvValue(v string) string { if v == "" { - return `""` + return `''` } - needsQuoting := strings.ContainsAny(v, " \t\"\\$`!#&|;(){}[]<>?*~'") + needsQuoting := strings.ContainsAny(v, shellQuotedSpecialChars) if !needsQuoting { return v } - var b strings.Builder - b.WriteByte('"') - for _, c := range v { - if c == '"' || c == '\\' || c == '$' || c == '`' || c == '!' { - b.WriteByte('\\') - } - b.WriteRune(c) - } - b.WriteByte('"') - return b.String() + return "'" + strings.ReplaceAll(v, "'", "'\\''") + "'" } diff --git a/cmd/auth/env_test.go b/cmd/auth/env_test.go index c094a33e25..66bc6a7e7f 100644 --- a/cmd/auth/env_test.go +++ b/cmd/auth/env_test.go @@ -4,6 +4,7 @@ import ( "bytes" "testing" + "github.com/databricks/cli/libs/cmdio" "github.com/databricks/cli/libs/flags" "github.com/spf13/cobra" "github.com/stretchr/testify/assert" @@ -17,15 +18,16 @@ func TestQuoteEnvValue(t *testing.T) { want string }{ {name: "simple value", in: "hello", want: "hello"}, - {name: "empty value", in: "", want: `""`}, - {name: "value with space", in: "hello world", want: `"hello world"`}, - {name: "value with tab", in: "hello\tworld", want: "\"hello\tworld\""}, - {name: "value with double quote", in: `say "hi"`, want: `"say \"hi\""`}, - {name: "value with backslash", in: `path\to`, want: `"path\\to"`}, + {name: "empty value", in: "", want: `''`}, + {name: "value with space", in: "hello world", want: "'hello world'"}, + {name: "value with tab", in: "hello\tworld", want: "'hello\tworld'"}, + {name: "value with double quote", in: `say "hi"`, want: "'say \"hi\"'"}, + {name: "value with backslash", in: `path\to`, want: "'path\\to'"}, {name: "url value", in: "https://example.com", want: "https://example.com"}, - {name: "value with dollar", in: "price$5", want: `"price\$5"`}, - {name: "value with backtick", in: "hello`world", want: `"hello\` + "`" + `world"`}, - {name: "value with bang", in: "hello!world", want: `"hello\!world"`}, + {name: "value with dollar", in: "price$5", want: "'price$5'"}, + {name: "value with backtick", in: "hello`world", want: "'hello`world'"}, + {name: "value with bang", in: "hello!world", want: "'hello!world'"}, + {name: "value with single quote", in: "it's", want: "'it'\\''s'"}, } for _, c := range cases { t.Run(c.name, func(t *testing.T) { @@ -51,6 +53,11 @@ func TestEnvCommand_TextOutput(t *testing.T) { args: []string{"--host", "https://test.cloud.databricks.com", "--output", "text"}, wantJSON: false, }, + { + name: "explicit --output json produces JSON", + args: []string{"--host", "https://test.cloud.databricks.com", "--output", "json"}, + wantJSON: true, + }, } for _, c := range cases { @@ -61,6 +68,7 @@ func TestEnvCommand_TextOutput(t *testing.T) { envCmd := newEnvCommand() parent.AddCommand(envCmd) + parent.SetContext(cmdio.MockDiscard(t.Context())) // Set DATABRICKS_TOKEN so the SDK's config.Authenticate succeeds // without hitting a real endpoint. diff --git a/cmd/auth/token.go b/cmd/auth/token.go index c8fb626f7d..9def38466d 100644 --- a/cmd/auth/token.go +++ b/cmd/auth/token.go @@ -85,22 +85,25 @@ using a client ID and secret is not supported.`, if err != nil { return err } + return writeTokenOutput(cmd, t) + } - // Output plain token when the user explicitly passes --output text. - if cmd.Flag("output").Changed && root.OutputType(cmd) == flags.OutputText { - _, _ = fmt.Fprintln(cmd.OutOrStdout(), t.AccessToken) - return nil - } + return cmd +} - raw, err := json.MarshalIndent(t, "", " ") - if err != nil { - return err - } - _, _ = cmd.OutOrStdout().Write(raw) +func writeTokenOutput(cmd *cobra.Command, t *oauth2.Token) error { + // Output plain token when the user explicitly passes --output text. + if cmd.Flag("output").Changed && root.OutputType(cmd) == flags.OutputText { + _, _ = fmt.Fprintln(cmd.OutOrStdout(), t.AccessToken) return nil } - return cmd + raw, err := json.MarshalIndent(t, "", " ") + if err != nil { + return err + } + _, _ = cmd.OutOrStdout().Write(raw) + return nil } type loadTokenArgs struct { diff --git a/cmd/auth/token_test.go b/cmd/auth/token_test.go index 289313bae8..0e41952995 100644 --- a/cmd/auth/token_test.go +++ b/cmd/auth/token_test.go @@ -3,13 +3,10 @@ package auth import ( "bytes" "context" - "encoding/json" - "fmt" "net/http" "testing" "time" - "github.com/databricks/cli/cmd/root" "github.com/databricks/cli/libs/auth" "github.com/databricks/cli/libs/cmdio" "github.com/databricks/cli/libs/databrickscfg/profile" @@ -796,9 +793,8 @@ func TestTokenCommand_TextOutput(t *testing.T) { parent.PersistentFlags().StringP("profile", "p", "", "~/.databrickscfg profile") tokenCmd := newTokenCommand(authArgs) - // Override RunE to inject test profiler and token cache. - // The output formatting logic below must mirror newTokenCommand.RunE. - // If you change the output logic in newTokenCommand, update this too. + // Override RunE to inject test profiler and token cache while reusing + // the production output formatter. tokenCmd.RunE = func(cmd *cobra.Command, args []string) error { profileName := "" if f := cmd.Flag("profile"); f != nil { @@ -815,16 +811,7 @@ func TestTokenCommand_TextOutput(t *testing.T) { if err != nil { return err } - if cmd.Flag("output").Changed && root.OutputType(cmd) == flags.OutputText { - _, _ = fmt.Fprintln(cmd.OutOrStdout(), tok.AccessToken) - return nil - } - raw, err := json.MarshalIndent(tok, "", " ") - if err != nil { - return err - } - _, _ = cmd.OutOrStdout().Write(raw) - return nil + return writeTokenOutput(cmd, tok) } parent.AddCommand(tokenCmd) From 72ebfd99d0b718b4a4f2f991a39a6196075924b3 Mon Sep 17 00:00:00 2001 From: simon Date: Fri, 13 Mar 2026 07:58:47 +0100 Subject: [PATCH 4/7] Fix gofmt formatting in quoteEnvValue --- cmd/auth/env.go | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/cmd/auth/env.go b/cmd/auth/env.go index d0ef95c474..2c8169a23e 100644 --- a/cmd/auth/env.go +++ b/cmd/auth/env.go @@ -173,7 +173,7 @@ func collectEnvVars(cfg *config.Config) map[string]string { // quoteEnvValue quotes a value for KEY=VALUE output if it contains spaces or // shell-special characters. Single quotes prevent shell expansion, and -// embedded single quotes use the POSIX-compatible '\'' sequence. +// embedded single quotes use the POSIX-compatible '\” sequence. func quoteEnvValue(v string) string { if v == "" { return `''` From 0c68a51251ea56960f4a4a42cb0187b4057a57e9 Mon Sep 17 00:00:00 2001 From: simon Date: Fri, 13 Mar 2026 15:39:38 +0100 Subject: [PATCH 5/7] Fix quoteEnvValue doc comment to match actual behavior The comment referenced the '\" escape sequence, but the code actually uses the POSIX '\'' sequence (end-quote, backslash-escaped literal single quote, re-open quote). --- cmd/auth/env.go | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/cmd/auth/env.go b/cmd/auth/env.go index 2c8169a23e..d0ef95c474 100644 --- a/cmd/auth/env.go +++ b/cmd/auth/env.go @@ -173,7 +173,7 @@ func collectEnvVars(cfg *config.Config) map[string]string { // quoteEnvValue quotes a value for KEY=VALUE output if it contains spaces or // shell-special characters. Single quotes prevent shell expansion, and -// embedded single quotes use the POSIX-compatible '\” sequence. +// embedded single quotes use the POSIX-compatible '\'' sequence. func quoteEnvValue(v string) string { if v == "" { return `''` From a6000ece3f3a3662a413c629e970bab8e190a005 Mon Sep 17 00:00:00 2001 From: simon Date: Fri, 13 Mar 2026 16:01:18 +0100 Subject: [PATCH 6/7] Fix quoteEnvValue to treat newline and carriage return as shell-special characters Values containing \n or \r were emitted unquoted, producing raw multi-line shell output instead of a single safe KEY=VALUE pair. Add both characters to shellQuotedSpecialChars so they trigger single-quoting. --- cmd/auth/env.go | 4 ++-- cmd/auth/env_test.go | 2 ++ 2 files changed, 4 insertions(+), 2 deletions(-) diff --git a/cmd/auth/env.go b/cmd/auth/env.go index d0ef95c474..0477c17c3f 100644 --- a/cmd/auth/env.go +++ b/cmd/auth/env.go @@ -32,7 +32,7 @@ func canonicalHost(host string) (string, error) { var ErrNoMatchingProfiles = errors.New("no matching profiles found") -const shellQuotedSpecialChars = " \t\"\\$`!#&|;(){}[]<>?*~'" +const shellQuotedSpecialChars = " \t\n\r\"\\$`!#&|;(){}[]<>?*~'" func resolveSection(cfg *config.Config, iniFile *config.File) (*ini.Section, error) { var candidates []*ini.Section @@ -173,7 +173,7 @@ func collectEnvVars(cfg *config.Config) map[string]string { // quoteEnvValue quotes a value for KEY=VALUE output if it contains spaces or // shell-special characters. Single quotes prevent shell expansion, and -// embedded single quotes use the POSIX-compatible '\'' sequence. +// embedded single quotes use the POSIX-compatible '\” sequence. func quoteEnvValue(v string) string { if v == "" { return `''` diff --git a/cmd/auth/env_test.go b/cmd/auth/env_test.go index 66bc6a7e7f..a95da721b2 100644 --- a/cmd/auth/env_test.go +++ b/cmd/auth/env_test.go @@ -28,6 +28,8 @@ func TestQuoteEnvValue(t *testing.T) { {name: "value with backtick", in: "hello`world", want: "'hello`world'"}, {name: "value with bang", in: "hello!world", want: "'hello!world'"}, {name: "value with single quote", in: "it's", want: "'it'\\''s'"}, + {name: "value with newline", in: "line1\nline2", want: "'line1\nline2'"}, + {name: "value with carriage return", in: "line1\rline2", want: "'line1\rline2'"}, } for _, c := range cases { t.Run(c.name, func(t *testing.T) { From fcfc26dcd5ccf5f8d201cf69a0e193130dacc4e0 Mon Sep 17 00:00:00 2001 From: simon Date: Sun, 22 Mar 2026 08:23:31 +0100 Subject: [PATCH 7/7] Refactor auth env to use CLI's resolved auth context Use root.MustAnyClient and auth.Env instead of custom profile/host resolution logic. This makes auth env return the environment variables for the exact identity the CLI is authenticated as, including bundle context and all standard auth resolution paths. Breaking changes: - Removed command-specific --host and --profile flags (use the inherited flags from the parent/root commands) - JSON output is a flat map instead of wrapped in {"env": ...} - Only the primary env var per attribute is emitted (via auth.Env) --- cmd/auth/env.go | 154 ++++++------------------------------------- cmd/auth/env_test.go | 19 +++--- 2 files changed, 30 insertions(+), 143 deletions(-) diff --git a/cmd/auth/env.go b/cmd/auth/env.go index 0477c17c3f..15e02cfce6 100644 --- a/cmd/auth/env.go +++ b/cmd/auth/env.go @@ -1,150 +1,48 @@ package auth import ( - "context" "encoding/json" - "errors" "fmt" - "io/fs" - "net/http" - "net/url" + "maps" + "slices" "strings" "github.com/databricks/cli/cmd/root" - "github.com/databricks/cli/libs/databrickscfg/profile" + "github.com/databricks/cli/libs/auth" + "github.com/databricks/cli/libs/cmdctx" "github.com/databricks/cli/libs/flags" - "github.com/databricks/databricks-sdk-go/config" "github.com/spf13/cobra" - "gopkg.in/ini.v1" ) -func canonicalHost(host string) (string, error) { - parsedHost, err := url.Parse(host) - if err != nil { - return "", err - } - // If the host is empty, assume the scheme wasn't included. - if parsedHost.Host == "" { - return "https://" + host, nil - } - return "https://" + parsedHost.Host, nil -} - -var ErrNoMatchingProfiles = errors.New("no matching profiles found") - -const shellQuotedSpecialChars = " \t\n\r\"\\$`!#&|;(){}[]<>?*~'" - -func resolveSection(cfg *config.Config, iniFile *config.File) (*ini.Section, error) { - var candidates []*ini.Section - configuredHost, err := canonicalHost(cfg.Host) - if err != nil { - return nil, err - } - for _, section := range iniFile.Sections() { - hash := section.KeysHash() - host, ok := hash["host"] - if !ok { - // if host is not set - continue - } - canonical, err := canonicalHost(host) - if err != nil { - // we're fine with other corrupt profiles - continue - } - if canonical != configuredHost { - continue - } - candidates = append(candidates, section) - } - if len(candidates) == 0 { - return nil, ErrNoMatchingProfiles - } - // in the real situations, we don't expect this to happen often - // (if not at all), hence we don't trim the list - if len(candidates) > 1 { - var profiles []string - for _, v := range candidates { - profiles = append(profiles, v.Name()) - } - return nil, fmt.Errorf("%s match %s in %s", - strings.Join(profiles, " and "), cfg.Host, cfg.ConfigFile) - } - return candidates[0], nil -} - -func loadFromDatabricksCfg(ctx context.Context, cfg *config.Config) error { - iniFile, err := profile.DefaultProfiler.Get(ctx) - if errors.Is(err, fs.ErrNotExist) { - // it's fine not to have ~/.databrickscfg - return nil - } - if err != nil { - return err - } - profile, err := resolveSection(cfg, iniFile) - if err == ErrNoMatchingProfiles { - // it's also fine for Azure CLI or Databricks CLI, which - // are resolved by unified auth handling in the Go SDK. - return nil - } - if err != nil { - return err - } - cfg.Profile = profile.Name() - return nil -} - func newEnvCommand() *cobra.Command { cmd := &cobra.Command{ Use: "env", - Short: "Get env", + Short: "Get authentication environment variables for the current CLI context", + Long: `Output the environment variables needed to authenticate as the same identity +the CLI is currently authenticated as. This is useful for configuring downstream +tools that accept Databricks authentication via environment variables.`, } - var host string - var profile string - cmd.Flags().StringVar(&host, "host", host, "Hostname to get auth env for") - cmd.Flags().StringVar(&profile, "profile", profile, "Profile to get auth env for") - cmd.RunE = func(cmd *cobra.Command, args []string) error { - cfg := &config.Config{ - Host: host, - Profile: profile, - } - if profile != "" { - cfg.Profile = profile - } else if cfg.Host == "" { - cfg.Profile = "DEFAULT" - } else if err := loadFromDatabricksCfg(cmd.Context(), cfg); err != nil { - return err - } - // Go SDK is lazy loaded because of Terraform semantics, - // so we're creating a dummy HTTP request as a placeholder - // for headers. - r := &http.Request{Header: http.Header{}} - err := cfg.Authenticate(r.WithContext(cmd.Context())) + _, err := root.MustAnyClient(cmd, args) if err != nil { return err } + + cfg := cmdctx.ConfigUsed(cmd.Context()) + envVars := auth.Env(cfg) + // Output KEY=VALUE lines when the user explicitly passes --output text. if cmd.Flag("output").Changed && root.OutputType(cmd) == flags.OutputText { w := cmd.OutOrStdout() - for _, a := range config.ConfigAttributes { - if a.IsZero(cfg) { - continue - } - v := a.GetString(cfg) - for _, envName := range a.EnvVars { - fmt.Fprintf(w, "%s=%s\n", envName, quoteEnvValue(v)) - } + keys := slices.Sorted(maps.Keys(envVars)) + for _, k := range keys { + fmt.Fprintf(w, "%s=%s\n", k, quoteEnvValue(envVars[k])) } return nil } - vars := collectEnvVars(cfg) - raw, err := json.MarshalIndent(map[string]any{ - "env": vars, - }, "", " ") + raw, err := json.MarshalIndent(envVars, "", " ") if err != nil { return err } @@ -155,25 +53,11 @@ func newEnvCommand() *cobra.Command { return cmd } -// collectEnvVars returns the environment variables for the given config -// as a map from env var name to value. -func collectEnvVars(cfg *config.Config) map[string]string { - vars := map[string]string{} - for _, a := range config.ConfigAttributes { - if a.IsZero(cfg) { - continue - } - v := a.GetString(cfg) - for _, envName := range a.EnvVars { - vars[envName] = v - } - } - return vars -} +const shellQuotedSpecialChars = " \t\n\r\"\\$`!#&|;(){}[]<>?*~'" // quoteEnvValue quotes a value for KEY=VALUE output if it contains spaces or // shell-special characters. Single quotes prevent shell expansion, and -// embedded single quotes use the POSIX-compatible '\” sequence. +// embedded single quotes use the POSIX-compatible '\" sequence. func quoteEnvValue(v string) string { if v == "" { return `''` diff --git a/cmd/auth/env_test.go b/cmd/auth/env_test.go index a95da721b2..d3347dd84d 100644 --- a/cmd/auth/env_test.go +++ b/cmd/auth/env_test.go @@ -47,35 +47,39 @@ func TestEnvCommand_TextOutput(t *testing.T) { }{ { name: "default output is JSON", - args: []string{"--host", "https://test.cloud.databricks.com"}, + args: nil, wantJSON: true, }, { name: "explicit --output text produces KEY=VALUE lines", - args: []string{"--host", "https://test.cloud.databricks.com", "--output", "text"}, + args: []string{"--output", "text"}, wantJSON: false, }, { name: "explicit --output json produces JSON", - args: []string{"--host", "https://test.cloud.databricks.com", "--output", "json"}, + args: []string{"--output", "json"}, wantJSON: true, }, } for _, c := range cases { t.Run(c.name, func(t *testing.T) { + // Isolate from real config/token cache on the machine. + t.Setenv("DATABRICKS_CONFIG_FILE", t.TempDir()+"/.databrickscfg") + t.Setenv("HOME", t.TempDir()) + // Set env vars so MustAnyClient resolves auth via PAT. + t.Setenv("DATABRICKS_HOST", "https://test.cloud.databricks.com") + t.Setenv("DATABRICKS_TOKEN", "test-token-value") + parent := &cobra.Command{Use: "databricks"} outputFlag := flags.OutputText parent.PersistentFlags().VarP(&outputFlag, "output", "o", "output type: text or json") + parent.PersistentFlags().StringP("profile", "p", "", "~/.databrickscfg profile") envCmd := newEnvCommand() parent.AddCommand(envCmd) parent.SetContext(cmdio.MockDiscard(t.Context())) - // Set DATABRICKS_TOKEN so the SDK's config.Authenticate succeeds - // without hitting a real endpoint. - t.Setenv("DATABRICKS_TOKEN", "test-token-value") - var buf bytes.Buffer parent.SetOut(&buf) parent.SetArgs(append([]string{"env"}, c.args...)) @@ -91,7 +95,6 @@ func TestEnvCommand_TextOutput(t *testing.T) { assert.NotContains(t, output, "{") assert.Contains(t, output, "DATABRICKS_HOST=") assert.Contains(t, output, "=") - // Verify KEY=VALUE format (no JSON structure) assert.NotContains(t, output, `"env"`) } })