From 3554e86ab70c16f8e9828210c83e58dff347686f Mon Sep 17 00:00:00 2001
From: Hector Castejon Diaz <hector.castejon@databricks.com>
Date: Wed, 19 Jun 2024 16:52:55 +0300
Subject: [PATCH] Fixes

---
 client/client.go             | 4 ++--
 httpclient/request.go        | 6 +-----
 openapi/code/method.go       | 2 +-
 openapi/code/service.go      | 6 +++---
 service/oauth2/data_plane.go | 8 ++++----
 5 files changed, 11 insertions(+), 15 deletions(-)

diff --git a/client/client.go b/client/client.go
index 77cd7e146..175a57351 100644
--- a/client/client.go
+++ b/client/client.go
@@ -47,12 +47,12 @@ func (c *DatabricksClient) ConfiguredAccountID() string {
 	return c.Config.AccountID
 }
 
-// Returns the inner Api Client.
+// ApiClient returns the inner Api Client.
 func (c *DatabricksClient) ApiClient() *httpclient.ApiClient {
 	return c.client
 }
 
-// Returns a new OAuth token using the provided token. The token must be a JWT token.
+// GetOAuthToken returns a new OAuth token using the provided token. The token must be a JWT token.
 // The resulting token is scoped to the authorization details provided.
 //
 // **NOTE:** Experimental: This API may change or be removed in a future release
diff --git a/httpclient/request.go b/httpclient/request.go
index 5bb6623d8..b8401fe19 100644
--- a/httpclient/request.go
+++ b/httpclient/request.go
@@ -35,11 +35,7 @@ func WithRequestHeaders(headers map[string]string) DoOption {
 
 // WithToken uses the specified golang.org/x/oauth2 token on a request
 func WithToken(token *oauth2.Token) DoOption {
-	visitor := WithRequestVisitor(func(r *http.Request) error {
-		auth := fmt.Sprintf("%s %s", token.TokenType, token.AccessToken)
-		r.Header.Set("Authorization", auth)
-		return nil
-	})
+	visitor := WithTokenSource(oauth2.StaticTokenSource(token))
 	visitor.isAuthOption = true
 	return visitor
 }
diff --git a/openapi/code/method.go b/openapi/code/method.go
index 7b3cee41a..df759724d 100644
--- a/openapi/code/method.go
+++ b/openapi/code/method.go
@@ -50,7 +50,7 @@ type Method struct {
 	shortcut   bool
 }
 
-// Returns the fields which contains the DataPlane info. Each field is nested in the previous one.
+// DataPlaneInfoFields returns the fields which contains the DataPlane info. Each field is nested in the previous one.
 func (m *Method) DataPlaneInfoFields() []*Field {
 	if m.DataPlane == nil {
 		return nil
diff --git a/openapi/code/service.go b/openapi/code/service.go
index 7e4af23f0..8f0b9441b 100644
--- a/openapi/code/service.go
+++ b/openapi/code/service.go
@@ -35,12 +35,12 @@ type Service struct {
 	tag                 *openapi.Tag
 }
 
-// Returns whether this API has a DataPlane counterpart.
+// HasDataPlaneAPI returns whether this API has a DataPlane counterpart.
 func (svc *Service) HasDataPlaneAPI() bool {
 	return len(svc.DataPlaneServices) > 0
 }
 
-// Returns the method in the Control Plane which contains the DataInfo object
+// DataPlaneInfoMethod returns the method in the Control Plane which contains the DataInfo object
 func (svc *Service) DataPlaneInfoMethod() *Method {
 	methodName := ""
 	for _, m := range svc.methods {
@@ -61,7 +61,7 @@ func (svc *Service) HasParent() bool {
 	return svc.tag.ParentService != ""
 }
 
-// Returns whether the service is a DataPlane service
+// IsDataPlane returns whether the service is a DataPlane service
 // This is determined by the service having a matching control plane service
 func (svc *Service) IsDataPlane() bool {
 	return svc.tag.ControlPlaneService != ""
diff --git a/service/oauth2/data_plane.go b/service/oauth2/data_plane.go
index fec3ef155..b090c08f0 100644
--- a/service/oauth2/data_plane.go
+++ b/service/oauth2/data_plane.go
@@ -6,11 +6,14 @@ import (
 	"golang.org/x/oauth2"
 )
 
+// DataPlaneHelper is a helper struct to fetch, refresh and manage DataPlane details and tokens.
 type DataPlaneHelper struct {
 	infos  map[string]*DataPlaneInfo
 	tokens map[string]*oauth2.Token
 }
 
+// GetDataPlaneDetails returns the endpoint URL and token. It returns a cached token if it is valid,
+// otherwise it refreshes the token and returns the new token.
 func (o *DataPlaneHelper) GetDataPlaneDetails(method string, params []string, refresh func(*DataPlaneInfo) (*oauth2.Token, error), infoGetter func() (*DataPlaneInfo, error)) (string, *oauth2.Token, error) {
 	if o.infos == nil {
 		o.infos = make(map[string]*DataPlaneInfo)
@@ -20,10 +23,6 @@ func (o *DataPlaneHelper) GetDataPlaneDetails(method string, params []string, re
 	}
 	key := o.generateKey(method, params)
 	info, infoOk := o.infos[key]
-	token, tokenOk := o.tokens[key]
-	if infoOk && tokenOk && token.Valid() {
-		return info.EndpointUrl, token, nil
-	}
 	if !infoOk {
 		newInfo, err := infoGetter()
 		if err != nil {
@@ -32,6 +31,7 @@ func (o *DataPlaneHelper) GetDataPlaneDetails(method string, params []string, re
 		o.infos[key] = newInfo
 		info = newInfo
 	}
+	token, tokenOk := o.tokens[key]
 	if !tokenOk || !token.Valid() {
 		newToken, err := refresh(info)
 		if err != nil {