From 77101c9b8567b17fd543f7c44cec7cd8c167bc68 Mon Sep 17 00:00:00 2001
From: hectorcast-db <hector.castejon@databricks.com>
Date: Wed, 11 Oct 2023 13:12:18 +0200
Subject: [PATCH] Use profile information when getting a token using the CLI
 (#855)

## Changes
Use stored profile information when the user provides the profile flag
when using the `databricks auth token` command.

## Tests
Run the command with and without the profile flag

```
./cli auth token
Databricks Host: https://e2-dogfood.staging.cloud.databricks.com/
{
  "access_token": "****",
  "token_type": "Bearer",
  "expiry": "2023-10-10T14:24:11.85617+02:00"
}%

./cli auth token --profile DEFAULT
{
  "access_token": "*****",
  "token_type": "Bearer",
  "expiry": "2023-10-10T14:24:11.85617+02:00"
}%

./cli auth token https://e2-dogfood.staging.cloud.databricks.com/
{
  "access_token": "*****",
  "token_type": "Bearer",
  "expiry": "2023-10-11T09:24:55.046029+02:00"
}%

./cli auth token --profile DEFAULT https://e2-dogfood.staging.cloud.databricks.com/
Error: providing both a profile and a host parameters is not supported
```
---
 cmd/auth/login.go | 30 +++++++++++++++++++-----------
 cmd/auth/token.go | 17 +++++++++++++++--
 2 files changed, 34 insertions(+), 13 deletions(-)

diff --git a/cmd/auth/login.go b/cmd/auth/login.go
index a14c5ebe..3a3f3a6d 100644
--- a/cmd/auth/login.go
+++ b/cmd/auth/login.go
@@ -60,20 +60,10 @@ func newLoginCommand(persistentAuth *auth.PersistentAuth) *cobra.Command {
 			profileName = profile
 		}
 
-		// If the chosen profile has a hostname and the user hasn't specified a host, infer the host from the profile.
-		_, profiles, err := databrickscfg.LoadProfiles(func(p databrickscfg.Profile) bool {
-			return p.Name == profileName
-		})
+		err := setHost(ctx, profileName, persistentAuth, args)
 		if err != nil {
 			return err
 		}
-		if persistentAuth.Host == "" {
-			if len(profiles) > 0 && profiles[0].Host != "" {
-				persistentAuth.Host = profiles[0].Host
-			} else {
-				configureHost(ctx, persistentAuth, args, 0)
-			}
-		}
 		defer persistentAuth.Close()
 
 		// We need the config without the profile before it's used to initialise new workspace client below.
@@ -135,3 +125,21 @@ func newLoginCommand(persistentAuth *auth.PersistentAuth) *cobra.Command {
 
 	return cmd
 }
+
+func setHost(ctx context.Context, profileName string, persistentAuth *auth.PersistentAuth, args []string) error {
+	// If the chosen profile has a hostname and the user hasn't specified a host, infer the host from the profile.
+	_, profiles, err := databrickscfg.LoadProfiles(func(p databrickscfg.Profile) bool {
+		return p.Name == profileName
+	})
+	if err != nil {
+		return err
+	}
+	if persistentAuth.Host == "" {
+		if len(profiles) > 0 && profiles[0].Host != "" {
+			persistentAuth.Host = profiles[0].Host
+		} else {
+			configureHost(ctx, persistentAuth, args, 0)
+		}
+	}
+	return nil
+}
diff --git a/cmd/auth/token.go b/cmd/auth/token.go
index 242a3dab..d763b956 100644
--- a/cmd/auth/token.go
+++ b/cmd/auth/token.go
@@ -3,6 +3,7 @@ package auth
 import (
 	"context"
 	"encoding/json"
+	"errors"
 	"time"
 
 	"github.com/databricks/cli/libs/auth"
@@ -21,8 +22,20 @@ func newTokenCommand(persistentAuth *auth.PersistentAuth) *cobra.Command {
 
 	cmd.RunE = func(cmd *cobra.Command, args []string) error {
 		ctx := cmd.Context()
-		if persistentAuth.Host == "" {
-			configureHost(ctx, persistentAuth, args, 0)
+
+		var profileName string
+		profileFlag := cmd.Flag("profile")
+		if profileFlag != nil {
+			profileName = profileFlag.Value.String()
+			// If a profile is provided we read the host from the .databrickscfg file
+			if profileName != "" && len(args) > 0 {
+				return errors.New("providing both a profile and a host parameters is not supported")
+			}
+		}
+
+		err := setHost(ctx, profileName, persistentAuth, args)
+		if err != nil {
+			return err
 		}
 		defer persistentAuth.Close()