diff --git a/cmd/auth/token.go b/cmd/auth/token.go index 06d5f2c2ac..9df3c3bb08 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" @@ -88,17 +90,29 @@ and secret is not supported.`, if err != nil { return err } - raw, err := json.MarshalIndent(t, "", " ") - if err != nil { - return err - } - _, _ = cmd.OutOrStdout().Write(raw) - return nil + return writeTokenOutput(cmd, t) } return cmd } +func writeTokenOutput(cmd *cobra.Command, t *oauth2.Token) error { + // Only honor the explicit --output text flag, not implicit text mode + // (e.g. from DATABRICKS_OUTPUT_FORMAT). auth token defaults to JSON, + // and changing that implicitly would break scripts that parse JSON output. + if cmd.Flag("output").Changed && root.OutputType(cmd) == flags.OutputText { + _, err := fmt.Fprintln(cmd.OutOrStdout(), t.AccessToken) + return err + } + + raw, err := json.MarshalIndent(t, "", " ") + if err != nil { + return err + } + _, err = cmd.OutOrStdout().Write(raw) + return err +} + type loadTokenArgs struct { // authArguments is the parsed auth arguments, including the host and optionally the account ID. authArguments *auth.AuthArguments diff --git a/cmd/auth/token_test.go b/cmd/auth/token_test.go index f36909c9b2..a522340549 100644 --- a/cmd/auth/token_test.go +++ b/cmd/auth/token_test.go @@ -1,6 +1,7 @@ package auth import ( + "bytes" "context" "errors" "net/http" @@ -11,8 +12,10 @@ import ( "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" ) @@ -798,3 +801,105 @@ 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 with newline", + args: []string{"--profile", "test-ws", "--output", "text"}, + wantSubstr: "new-access-token\n", + 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 reusing + // the production output formatter. + 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 + } + return writeTokenOutput(cmd, tok) + } + + 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, "{") + } + }) + } +}