From 46f61fad09598fe4ae7f18f4c3bd59d8c469a827 Mon Sep 17 00:00:00 2001
From: Miles Yucht <miles@databricks.com>
Date: Mon, 11 Mar 2024 21:58:07 +0100
Subject: [PATCH] WIP: DATABRICKS_CA_BUNDLE

---
 cmd/auth/login.go  |  8 ++++----
 cmd/auth/token.go  |  6 +++---
 go.mod             |  2 ++
 libs/auth/oauth.go | 28 ++++++++++++++++++++++++++--
 4 files changed, 35 insertions(+), 9 deletions(-)

diff --git a/cmd/auth/login.go b/cmd/auth/login.go
index b0bc7a853b..cebc6de380 100644
--- a/cmd/auth/login.go
+++ b/cmd/auth/login.go
@@ -60,10 +60,10 @@ func newLoginCommand(persistentAuth *auth.PersistentAuth) *cobra.Command {
 			if err != nil {
 				return err
 			}
-			profileName = profile
+			persistentAuth.Profile = profile
 		}
 
-		err := setHost(ctx, profileName, persistentAuth, args)
+		err := setHost(ctx, persistentAuth, args)
 		if err != nil {
 			return err
 		}
@@ -127,10 +127,10 @@ func newLoginCommand(persistentAuth *auth.PersistentAuth) *cobra.Command {
 	return cmd
 }
 
-func setHost(ctx context.Context, profileName string, persistentAuth *auth.PersistentAuth, args []string) error {
+func setHost(ctx context.Context, 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(ctx, func(p databrickscfg.Profile) bool {
-		return p.Name == profileName
+		return p.Name == persistentAuth.Profile
 	})
 	// Tolerate ErrNoConfiguration here, as we will write out a configuration as part of the login flow.
 	if err != nil && !errors.Is(err, databrickscfg.ErrNoConfiguration) {
diff --git a/cmd/auth/token.go b/cmd/auth/token.go
index d763b95642..2f6635c03e 100644
--- a/cmd/auth/token.go
+++ b/cmd/auth/token.go
@@ -26,14 +26,14 @@ func newTokenCommand(persistentAuth *auth.PersistentAuth) *cobra.Command {
 		var profileName string
 		profileFlag := cmd.Flag("profile")
 		if profileFlag != nil {
-			profileName = profileFlag.Value.String()
+			persistentAuth.Profile = 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")
+				return errors.New("providing both a profile and a hostname is not supported")
 			}
 		}
 
-		err := setHost(ctx, profileName, persistentAuth, args)
+		err := setHost(ctx, persistentAuth, args)
 		if err != nil {
 			return err
 		}
diff --git a/go.mod b/go.mod
index 832efbc664..2e59019622 100644
--- a/go.mod
+++ b/go.mod
@@ -69,3 +69,5 @@ require (
 	google.golang.org/protobuf v1.32.0 // indirect
 	gopkg.in/yaml.v2 v2.4.0 // indirect
 )
+
+replace github.com/databricks/databricks-sdk-go => /Users/miles/databricks-sdk-go
diff --git a/libs/auth/oauth.go b/libs/auth/oauth.go
index dd27d04b2e..6cdae7d738 100644
--- a/libs/auth/oauth.go
+++ b/libs/auth/oauth.go
@@ -16,6 +16,8 @@ import (
 	"time"
 
 	"github.com/databricks/cli/libs/auth/cache"
+	"github.com/databricks/databricks-sdk-go/config"
+	"github.com/databricks/databricks-sdk-go/httpclient"
 	"github.com/databricks/databricks-sdk-go/retries"
 	"github.com/pkg/browser"
 	"golang.org/x/oauth2"
@@ -42,6 +44,7 @@ var ( // Databricks SDK API: `databricks OAuth is not` will be checked for prese
 type PersistentAuth struct {
 	Host      string
 	AccountID string
+	Profile   string
 
 	http    httpGet
 	cache   tokenCache
@@ -82,6 +85,7 @@ func (a *PersistentAuth) Load(ctx context.Context) (*oauth2.Token, error) {
 		return nil, err
 	}
 	// eagerly refresh token
+	ctx = context.WithValue(ctx, oauth2.HTTPClient, a.http)
 	refreshed, err := cfg.TokenSource(ctx, t).Token()
 	if err != nil {
 		return nil, fmt.Errorf("token refresh: %w", err)
@@ -96,7 +100,9 @@ func (a *PersistentAuth) Load(ctx context.Context) (*oauth2.Token, error) {
 }
 
 func (a *PersistentAuth) ProfileName() string {
-	// TODO: get profile name from interactive input
+	if a.Profile != "" {
+		return a.Profile
+	}
 	if a.AccountID != "" {
 		return fmt.Sprintf("ACCOUNT-%s", a.AccountID)
 	}
@@ -138,7 +144,25 @@ func (a *PersistentAuth) init(ctx context.Context) error {
 		return ErrFetchCredentials
 	}
 	if a.http == nil {
-		a.http = http.DefaultClient
+		c := &config.Config{
+			Profile:   a.Profile,
+			Host:      a.Host,
+			AccountID: a.AccountID,
+		}
+		c.EnsureResolved()
+		clientConfig := httpclient.ClientConfig{
+			DebugHeaders:       c.DebugHeaders,
+			DebugTruncateBytes: c.DebugTruncateBytes,
+			InsecureSkipVerify: c.InsecureSkipVerify,
+			CABundle:           c.CABundle,
+			RetryTimeout:       time.Duration(c.RetryTimeoutSeconds) * time.Second,
+			HTTPTimeout:        time.Duration(c.HTTPTimeoutSeconds) * time.Second,
+		}
+		httpClient, err := httpclient.NewHttpClient(clientConfig)
+		if err != nil {
+			return err
+		}
+		a.http = httpClient
 	}
 	if a.cache == nil {
 		a.cache = &cache.TokenCache{}