Merge pull request #78 from databrickslabs/fix-azure-workspace-creation-bug

Fix azure workspace creation bug

Author: stikkireddy

client/service/apis.go

@@ -17,22 +17,22 @@ func (c *DBApiClient) SetConfig(clientConfig *DBApiClientConfig) DBApiClient {
 }
 
 // Clusters returns an instance of ClustersAPI
-func (c DBApiClient) Clusters() ClustersAPI {
+func (c *DBApiClient) Clusters() ClustersAPI {
 	return ClustersAPI{Client: c}
 }
 
 // Secrets returns an instance of SecretsAPI
-func (c DBApiClient) Secrets() SecretsAPI {
+func (c *DBApiClient) Secrets() SecretsAPI {
 	return SecretsAPI{Client: c}
 }
 
 // SecretScopes returns an instance of SecretScopesAPI
-func (c DBApiClient) SecretScopes() SecretScopesAPI {
+func (c *DBApiClient) SecretScopes() SecretScopesAPI {
 	return SecretScopesAPI{Client: c}
 }
 
 // SecretAcls returns an instance of SecretAclsAPI
-func (c DBApiClient) SecretAcls() SecretAclsAPI {
+func (c *DBApiClient) SecretAcls() SecretAclsAPI {
 	return SecretAclsAPI{Client: c}
 }
 
@@ -42,50 +42,54 @@ func (c *DBApiClient) Tokens() TokensAPI {
 }
 
 // Users returns an instance of UsersAPI
-func (c DBApiClient) Users() UsersAPI {
+func (c *DBApiClient) Users() UsersAPI {
 	return UsersAPI{Client: c}
 }
 
 // Groups returns an instance of GroupsAPI
-func (c DBApiClient) Groups() GroupsAPI {
+func (c *DBApiClient) Groups() GroupsAPI {
 	return GroupsAPI{Client: c}
 }
 
 // Notebooks returns an instance of NotebooksAPI
-func (c DBApiClient) Notebooks() NotebooksAPI {
+func (c *DBApiClient) Notebooks() NotebooksAPI {
 	return NotebooksAPI{Client: c}
 }
 
 // Jobs returns an instance of JobsAPI
-func (c DBApiClient) Jobs() JobsAPI {
+func (c *DBApiClient) Jobs() JobsAPI {
 	return JobsAPI{Client: c}
 }
 
 // DBFS returns an instance of DBFSAPI
-func (c DBApiClient) DBFS() DBFSAPI {
+func (c *DBApiClient) DBFS() DBFSAPI {
 	return DBFSAPI{Client: c}
 }
 
 // Libraries returns an instance of LibrariesAPI
-func (c DBApiClient) Libraries() LibrariesAPI {
+func (c *DBApiClient) Libraries() LibrariesAPI {
 	return LibrariesAPI{Client: c}
 }
 
 // InstancePools returns an instance of InstancePoolsAPI
-func (c DBApiClient) InstancePools() InstancePoolsAPI {
+func (c *DBApiClient) InstancePools() InstancePoolsAPI {
 	return InstancePoolsAPI{Client: c}
 }
 
 // InstanceProfiles returns an instance of InstanceProfilesAPI
-func (c DBApiClient) InstanceProfiles() InstanceProfilesAPI {
+func (c *DBApiClient) InstanceProfiles() InstanceProfilesAPI {
 	return InstanceProfilesAPI{Client: c}
 }
 
 // Commands returns an instance of CommandsAPI
-func (c DBApiClient) Commands() CommandsAPI {
+func (c *DBApiClient) Commands() CommandsAPI {
 	return CommandsAPI{Client: c}
 }
 
-func (c DBApiClient) performQuery(method, path string, apiVersion string, headers map[string]string, data interface{}, secretsMask *SecretsMask) ([]byte, error) {
+func (c *DBApiClient) performQuery(method, path string, apiVersion string, headers map[string]string, data interface{}, secretsMask *SecretsMask) ([]byte, error) {
+	err := c.Config.getOrCreateToken()
+	if err != nil {
+		return []byte(""), err
+	}
 	return PerformQuery(c.Config, method, path, apiVersion, headers, true, false, data, secretsMask)
 }

client/service/client.go

@@ -12,6 +12,7 @@ import (
 	"net/url"
 	"reflect"
 	"strings"
+	"sync"
 	"time"
 
 	"github.com/google/go-querystring/query"
@@ -56,6 +57,8 @@ const (
 	BasicAuth AuthType = "BASIC"
 )
 
+var clientAuthorizerMutex sync.Mutex
+
 // DBApiClientConfig is used to configure the DataBricks Client
 type DBApiClientConfig struct {
 	Host               string
@@ -65,12 +68,14 @@ type DBApiClientConfig struct {
 	DefaultHeaders     map[string]string
 	InsecureSkipVerify bool
 	TimeoutSeconds     int
+	CustomAuthorizer   func(*DBApiClientConfig) error
 	client             *retryablehttp.Client
 }
 
 var transientErrorStringMatches []string = []string{ // TODO: Should we make these regexes to match more of the message or is this sufficient?
 	"com.databricks.backend.manager.util.UnknownWorkerEnvironmentException",
 	"does not have any associated worker environments",
+	"There is no worker environment with id",
 }
 
 // Setup initializes the client
@@ -123,7 +128,6 @@ func checkHTTPRetry(ctx context.Context, resp *http.Response, err error) (bool,
 		if err != nil {
 			return false, err
 		}
-
 		var errorBody DBApiErrorBody
 		err = json.Unmarshal(body, &errorBody)
 		if err != nil {
@@ -145,7 +149,22 @@ func checkHTTPRetry(ctx context.Context, resp *http.Response, err error) (bool,
 	return false, nil
 }
 
-func (c *DBApiClientConfig) getAuthHeader() map[string]string {
+func (c *DBApiClientConfig) getOrCreateToken() error {
+	if c.CustomAuthorizer != nil {
+		// Lock incase terraform tries to getOrCreateToken from multiple go routines on the same client ptr.
+		clientAuthorizerMutex.Lock()
+		defer clientAuthorizerMutex.Unlock()
+		if reflect.ValueOf(c.Token).IsZero() {
+			log.Println("NOT AUTHORIZED SO ATTEMPTING TO AUTHORIZE")
+			return c.CustomAuthorizer(c)
+		}
+		log.Println("ALREADY AUTHORIZED")
+		return nil
+	}
+	return nil
+}
+
+func (c DBApiClientConfig) getAuthHeader() map[string]string {
 	auth := make(map[string]string)
 	if c.AuthType == BasicAuth {
 		auth["Authorization"] = "Basic " + c.Token

client/service/clusters.go

@@ -11,7 +11,7 @@ import (
 
 // ClustersAPI is a struct that contains the Databricks api client to perform queries
 type ClustersAPI struct {
-	Client DBApiClient
+	Client *DBApiClient
 }
 
 // Create creates a new Spark cluster

client/service/commands.go

@@ -12,7 +12,7 @@ import (
 
 // CommandsAPI exposes the Context & Commands API
 type CommandsAPI struct {
-	Client DBApiClient
+	Client *DBApiClient
 }
 
 // Execute creates a spark context and executes a command and then closes context

client/service/dbfs.go

@@ -10,7 +10,7 @@ import (
 
 // DBFSAPI exposes the DBFS API
 type DBFSAPI struct {
-	Client DBApiClient
+	Client *DBApiClient
 }
 
 // Create creates a file in DBFS given data string in base64

client/service/groups.go

@@ -12,7 +12,7 @@ import (
 
 // GroupsAPI exposes the scim groups API
 type GroupsAPI struct {
-	Client DBApiClient
+	Client *DBApiClient
 }
 
 // Create creates a scim group in the Databricks workspace

client/service/instance_pools.go

@@ -8,7 +8,7 @@ import (
 
 // InstancePoolsAPI exposes the instance pools api
 type InstancePoolsAPI struct {
-	Client DBApiClient
+	Client *DBApiClient
 }
 
 // Create creates the instance pool to given the instance pool configuration

client/service/instance_profiles.go

@@ -9,7 +9,7 @@ import (
 
 // InstanceProfilesAPI exposes the instance profiles api on the AWS deployment of Databricks
 type InstanceProfilesAPI struct {
-	Client DBApiClient
+	Client *DBApiClient
 }
 
 // Create creates an instance profile record on Databricks

client/service/jobs.go

@@ -8,7 +8,7 @@ import (
 
 // JobsAPI exposes the Jobs API
 type JobsAPI struct {
-	Client DBApiClient
+	Client *DBApiClient
 }
 
 // Create creates a job on the workspace given the job settings

client/service/libraries.go

@@ -8,7 +8,7 @@ import (
 
 // LibrariesAPI exposes the Library API
 type LibrariesAPI struct {
-	Client DBApiClient
+	Client *DBApiClient
 }
 
 // Create installs the list of libraries given a cluster id