Add text output mode for auth token

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

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 "<error>", 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, "{")
+			}
+		})
+	}
+}