Use new workspace_url format + fix azure test (#114)

* Use new workspace_url format + fix azure test Co-authored-by: [email protected] <[email protected]> * Fix tests * Make logging optional * Remove comment * Add spaces Like a maniac * Fix fmt and fmt code * Add local install make command * Make local install use repo build Co-authored-by: Lawrence Gripper <[email protected]> * Remove redudant azure_auth * Get workspace url from azure managment api * Update comment * Change from pointers Co-authored-by: Stuart Leeks <[email protected]> * Update docs * Assert env var is set * Reverse bool * Use correct operator * Run all integration tests * Use KEY Co-authored-by: Lawrence Gripper <[email protected]> Co-authored-by: stikkireddy <[email protected]> Co-authored-by: Stuart Leeks <[email protected]>
databricks · Jun 24, 2020 · c23a1a6 · c23a1a6
1 parent 0686299
commit c23a1a6
Show file tree

Hide file tree

Showing 13 changed files with 200 additions and 209 deletions.
diff --git a/.devcontainer/Dockerfile b/.devcontainer/Dockerfile
@@ -37,7 +37,9 @@ RUN \
     # --> Delve for debugging
     go get github.com/go-delve/delve/cmd/[email protected] \
     # --> Go language server
-    && go get golang.org/x/tools/[email protected] \
+    && go get golang.org/x/tools/[email protected] \
+    # --> Goimports
+    && go get golang.org/x/tools/cmd/goimports \
     # --> Gotestsum
     && go get gotest.tools/[email protected] \
     # --> Go symbols and outline for go to symbol support and test support 

diff --git a/.devcontainer/devcontainer.json b/.devcontainer/devcontainer.json
@@ -3,6 +3,14 @@
 {
 	"name": "databricks-terraform",
 	"dockerFile": "Dockerfile",
+	"mounts": [
+		// Mount go mod cache
+		"source=databricks-terraform-gomodcache,target=/go/pkg,type=volume",
+		// Keep command history 
+		"source=databricks-terraform-bashhistory,target=/commandhistory,type=volume",
+		// Mount docker socket for docker builds
+		"source=/var/run/docker.sock,target=/var/run/docker.sock,type=bind",
+	],
 	"runArgs": [
 		// Uncomment the next line to use a non-root user. On Linux, this will prevent
 		// new files getting created as root, but you may need to update the USER_UID
@@ -12,14 +20,8 @@
 		"--security-opt",
 		"seccomp=unconfined",
 		"--privileged",
-
-		"--name", "databricks-terraform-devcontainer",
-		// Mount go mod cache
-		"-v", "databricks-terraform-gomodcache:/go/pkg",
-		// Keep command history 
-		"-v", "databricks-terraform-bashhistory:/root/commandhistory",
-		// Mount docker socket for docker builds
-		"-v", "/var/run/docker.sock:/var/run/docker.sock",
+		"--name",
+		"databricks-terraform-devcontainer",
 		// Use host network
 		"--network=host",
 	],
@@ -51,8 +53,9 @@
 	// "postCreateCommand": "go version",
 	// Add the IDs of extensions you want installed when the container is created in the array below.
 	"extensions": [
-        "ms-vscode.go",
+		"golang.go",
 		"mauve.terraform",
+		"hashicorp.terraform",
 		"ms-vsliveshare.vsliveshare"
 	]
 }
diff --git a/Makefile b/Makefile
@@ -48,10 +48,13 @@ vendor:
 	@echo "==> Filling vendor folder with library code..."
 	@go mod vendor
 
+local-install: build
+	mv terraform-provider-databricks $(HOME)/.terraform.d/plugins/terraform-provider-databricks_v0.2.0
+
 # INTEGRATION TESTING WITH AZURE
 terraform-acc-azure: lint
 	@echo "==> Running Terraform Acceptance Tests for Azure..."
-	@CLOUD_ENV="azure" TF_ACC=1 gotestsum --format short-verbose --raw-command go test -v -json -tags=azure  -short -coverprofile=coverage.out ./...
+	@/bin/bash integration-environment-azure/run.sh
 
 # INTEGRATION TESTING WITH AWS
 terraform-acc-aws: lint

diff --git a/databricks/azure_auth.go b/databricks/azure_auth.go
@@ -5,7 +5,7 @@ import (
 	"fmt"
 	"log"
 	"net/http"
-	urlParse "net/url"
+	"time"
 
 	"github.com/Azure/go-autorest/autorest/adal"
 	"github.com/Azure/go-autorest/autorest/azure"
@@ -17,15 +17,6 @@ const (
 	ADBResourceID string = "2ff814a6-3304-4ab8-85cb-cd0e6f879c1d"
 )
 
-// AzureAuth is a struct that contains information about the azure sp authentication
-type AzureAuth struct {
-	TokenPayload           *TokenPayload
-	ManagementToken        string
-	AdbWorkspaceResourceID string
-	AdbAccessToken         string
-	AdbPlatformToken       string
-}
-
 // TokenPayload contains all the auth information for azure sp authentication
 type TokenPayload struct {
 	ManagedResourceGroup string
@@ -38,6 +29,39 @@ type TokenPayload struct {
 	TenantID             string
 }
 
+type Workspace struct {
+	Name string `json:"name"`
+	ID   string `json:"id"`
+	Type string `json:"type"`
+	Sku  struct {
+		Name string `json:"name"`
+	} `json:"sku"`
+	Location   string `json:"location"`
+	Properties struct {
+		ManagedResourceGroupID string      `json:"managedResourceGroupId"`
+		Parameters             interface{} `json:"parameters"`
+		ProvisioningState      string      `json:"provisioningState"`
+		UIDefinitionURI        string      `json:"uiDefinitionUri"`
+		Authorizations         []struct {
+			PrincipalID      string `json:"principalId"`
+			RoleDefinitionID string `json:"roleDefinitionId"`
+		} `json:"authorizations"`
+		CreatedBy struct {
+			Oid           string `json:"oid"`
+			Puid          string `json:"puid"`
+			ApplicationID string `json:"applicationId"`
+		} `json:"createdBy"`
+		UpdatedBy struct {
+			Oid           string `json:"oid"`
+			Puid          string `json:"puid"`
+			ApplicationID string `json:"applicationId"`
+		} `json:"updatedBy"`
+		CreatedDateTime time.Time `json:"createdDateTime"`
+		WorkspaceID     string    `json:"workspaceId"`
+		WorkspaceURL    string    `json:"workspaceUrl"`
+	} `json:"properties"`
+}
+
 // WsProps contains information about the workspace properties
 type WsProps struct {
 	ManagedResourceGroupID string `json:"managedResourceGroupId"`
@@ -50,92 +74,90 @@ type WorkspaceRequest struct {
 	Location   string   `json:"location"`
 }
 
-func (a *AzureAuth) getManagementToken() error {
+func (t *TokenPayload) getManagementToken() (string, error) {
 	log.Println("[DEBUG] Creating Azure Databricks management OAuth token.")
 	mgmtTokenOAuthCfg, err := adal.NewOAuthConfigWithAPIVersion(azure.PublicCloud.ActiveDirectoryEndpoint,
-		a.TokenPayload.TenantID,
+		t.TenantID,
 		nil)
 	if err != nil {
-		return err
+		return "", err
 	}
 	mgmtToken, err := adal.NewServicePrincipalToken(
 		*mgmtTokenOAuthCfg,
-		a.TokenPayload.ClientID,
-		a.TokenPayload.ClientSecret,
+		t.ClientID,
+		t.ClientSecret,
 		azure.PublicCloud.ServiceManagementEndpoint)
 	if err != nil {
-		return err
+		return "", err
 	}
 	err = mgmtToken.Refresh()
 	if err != nil {
-		return err
+		return "", err
 	}
-	a.ManagementToken = mgmtToken.OAuthToken()
-	return nil
+
+	return mgmtToken.OAuthToken(), nil
 }
 
-func (a *AzureAuth) getWorkspaceID(config *service.DBApiClientConfig) error {
+func (t *TokenPayload) getWorkspace(config *service.DBApiClientConfig, managementToken string) (*Workspace, error) {
 	log.Println("[DEBUG] Getting Workspace ID via management token.")
-	// Escape all the ids
 	url := fmt.Sprintf("https://management.azure.com/subscriptions/%s/resourceGroups/%s"+
 		"/providers/Microsoft.Databricks/workspaces/%s",
-		urlParse.PathEscape(a.TokenPayload.SubscriptionID),
-		urlParse.PathEscape(a.TokenPayload.ResourceGroup),
-		urlParse.PathEscape(a.TokenPayload.WorkspaceName))
+		t.SubscriptionID,
+		t.ResourceGroup,
+		t.WorkspaceName)
 	headers := map[string]string{
 		"Content-Type":  "application/json",
 		"cache-control": "no-cache",
-		"Authorization": "Bearer " + a.ManagementToken,
+		"Authorization": "Bearer " + managementToken,
 	}
 	type apiVersion struct {
 		APIVersion string `url:"api-version"`
 	}
 	uriPayload := apiVersion{
 		APIVersion: "2018-04-01",
 	}
-	var responseMap map[string]interface{}
 	resp, err := service.PerformQuery(config, http.MethodGet, url, "2.0", headers, false, true, uriPayload, nil)
 	if err != nil {
-		return err
+		return nil, err
 	}
-	err = json.Unmarshal(resp, &responseMap)
+
+	var workspace Workspace
+	err = json.Unmarshal(resp, &workspace)
 	if err != nil {
-		return err
+		return nil, err
 	}
-	a.AdbWorkspaceResourceID = responseMap["id"].(string)
-	return err
+	return &workspace, err
 }
 
-func (a *AzureAuth) getADBPlatformToken() error {
+func (t *TokenPayload) getADBPlatformToken() (string, error) {
 	log.Println("[DEBUG] Creating Azure Databricks management OAuth token.")
 	platformTokenOAuthCfg, err := adal.NewOAuthConfigWithAPIVersion(azure.PublicCloud.ActiveDirectoryEndpoint,
-		a.TokenPayload.TenantID,
+		t.TenantID,
 		nil)
 	if err != nil {
-		return err
+		return "", err
 	}
 	platformToken, err := adal.NewServicePrincipalToken(
 		*platformTokenOAuthCfg,
-		a.TokenPayload.ClientID,
-		a.TokenPayload.ClientSecret,
+		t.ClientID,
+		t.ClientSecret,
 		ADBResourceID)
 	if err != nil {
-		return err
+		return "", err
 	}
 
 	err = platformToken.Refresh()
 	if err != nil {
-		return err
+		return "", err
 	}
-	a.AdbPlatformToken = platformToken.OAuthToken()
-	return nil
+
+	return platformToken.OAuthToken(), nil
 }
 
-func (a *AzureAuth) getWorkspaceAccessToken(config *service.DBApiClientConfig) error {
+func getWorkspaceAccessToken(config *service.DBApiClientConfig, managementToken, adbWorkspaceUrl, adbWorkspaceResourceID, adbPlatformToken string) (string, error) {
 	log.Println("[DEBUG] Creating workspace token")
-	apiLifeTimeInSeconds := int32(600)
+	apiLifeTimeInSeconds := int32(3600)
 	comment := "Secret made via SP"
-	url := "https://" + a.TokenPayload.AzureRegion + ".azuredatabricks.net/api/2.0/token/create"
 	payload := struct {
 		LifetimeSeconds int32  `json:"lifetime_seconds,omitempty"`
 		Comment         string `json:"comment,omitempty"`
@@ -145,60 +167,60 @@ func (a *AzureAuth) getWorkspaceAccessToken(config *service.DBApiClientConfig) e
 	}
 	headers := map[string]string{
 		"Content-Type": "application/x-www-form-urlencoded",
-		"X-Databricks-Azure-Workspace-Resource-Id": a.AdbWorkspaceResourceID,
-		"X-Databricks-Azure-SP-Management-Token":   a.ManagementToken,
+		"X-Databricks-Azure-Workspace-Resource-Id": adbWorkspaceResourceID,
+		"X-Databricks-Azure-SP-Management-Token":   managementToken,
 		"cache-control":                            "no-cache",
-		"Authorization":                            "Bearer " + a.AdbPlatformToken,
+		"Authorization":                            "Bearer " + adbPlatformToken,
 	}
 
-	var responseMap map[string]interface{}
+	url := adbWorkspaceUrl + "/api/2.0/token/create"
 	resp, err := service.PerformQuery(config, http.MethodPost, url, "2.0", headers, true, true, payload, nil)
 	if err != nil {
-		return err
+		return "", err
 	}
+
+	var responseMap map[string]interface{}
 	err = json.Unmarshal(resp, &responseMap)
 	if err != nil {
-		return err
+		return "", err
 	}
-	a.AdbAccessToken = responseMap["token_value"].(string)
-	return nil
+
+	return responseMap["token_value"].(string), nil
 }
 
 // Main function call that gets made and it follows 4 steps at the moment:
 // 1. Get Management OAuth Token using management endpoint
-// 2. Get Workspace ID
+// 2. Get Workspace ID and URL
 // 3. Get Azure Databricks Platform OAuth Token using Databricks resource id
 // 4. Get Azure Databricks Workspace Personal Access Token for the SP (60 min duration)
-func (a *AzureAuth) initWorkspaceAndGetClient(config *service.DBApiClientConfig) error {
-	//var dbClient service.DBApiClient
-
+func (t *TokenPayload) initWorkspaceAndGetClient(config *service.DBApiClientConfig) error {
 	// Get management token
-	err := a.getManagementToken()
+	managementToken, err := t.getManagementToken()
 	if err != nil {
 		return err
 	}
 
 	// Get workspace access token
-	err = a.getWorkspaceID(config)
+	adbWorkspace, err := t.getWorkspace(config, managementToken)
 	if err != nil {
 		return err
 	}
+	adbWorkspaceUrl := "https://" + adbWorkspace.Properties.WorkspaceURL
 
 	// Get platform token
-	err = a.getADBPlatformToken()
+	adbPlatformToken, err := t.getADBPlatformToken()
 	if err != nil {
 		return err
 	}
 
 	// Get workspace personal access token
-	err = a.getWorkspaceAccessToken(config)
+	workspaceAccessToken, err := getWorkspaceAccessToken(config, managementToken, adbWorkspaceUrl, adbWorkspace.ID, adbPlatformToken)
 	if err != nil {
 		return err
 	}
 
-	//// TODO: Eventually change this to include new Databricks domain names. May have to add new vars and/or deprecate existing args.
-	config.Host = "https://" + a.TokenPayload.AzureRegion + ".azuredatabricks.net"
-	config.Token = a.AdbAccessToken
+	config.Host = adbWorkspaceUrl
+	config.Token = workspaceAccessToken
 
 	return nil
 }
diff --git a/databricks/azure_auth_test.go b/databricks/azure_auth_test.go
@@ -1,6 +1,7 @@
 package databricks
 
 import (
+	"fmt"
 	"os"
 	"testing"
 
@@ -20,24 +21,19 @@ func TestAzureAuthCreateApiToken(t *testing.T) {
 		t.Skip("skipping integration test in short mode.")
 	}
 
-	azureAuth := AzureAuth{
-		TokenPayload: &TokenPayload{
-			ManagedResourceGroup: os.Getenv("DATABRICKS_AZURE_MANAGED_RESOURCE_GROUP"),
-			AzureRegion:          os.Getenv("AZURE_REGION"),
-			WorkspaceName:        os.Getenv("DATABRICKS_AZURE_WORKSPACE_NAME"),
-			ResourceGroup:        os.Getenv("DATABRICKS_AZURE_RESOURCE_GROUP"),
-		},
-		ManagementToken:        "",
-		AdbWorkspaceResourceID: "",
-		AdbAccessToken:         "",
-		AdbPlatformToken:       "",
+	tokenPayload := TokenPayload{
+		ManagedResourceGroup: getAndAssertEnv(t, "DATABRICKS_AZURE_MANAGED_RESOURCE_GROUP"),
+		AzureRegion:          getAndAssertEnv(t, "AZURE_REGION"),
+		WorkspaceName:        getAndAssertEnv(t, "DATABRICKS_AZURE_WORKSPACE_NAME"),
+		ResourceGroup:        getAndAssertEnv(t, "DATABRICKS_AZURE_RESOURCE_GROUP"),
+		SubscriptionID:       getAndAssertEnv(t, "DATABRICKS_AZURE_SUBSCRIPTION_ID"),
+		TenantID:             getAndAssertEnv(t, "DATABRICKS_AZURE_TENANT_ID"),
+		ClientID:             getAndAssertEnv(t, "DATABRICKS_AZURE_CLIENT_ID"),
+		ClientSecret:         getAndAssertEnv(t, "DATABRICKS_AZURE_CLIENT_SECRET"),
 	}
-	azureAuth.TokenPayload.SubscriptionID = os.Getenv("DATABRICKS_AZURE_SUBSCRIPTION_ID")
-	azureAuth.TokenPayload.TenantID = os.Getenv("DATABRICKS_AZURE_TENANT_ID")
-	azureAuth.TokenPayload.ClientID = os.Getenv("DATABRICKS_AZURE_CLIENT_ID")
-	azureAuth.TokenPayload.ClientSecret = os.Getenv("DATABRICKS_AZURE_CLIENT_SECRET")
+
 	config := GetIntegrationDBClientOptions()
-	err := azureAuth.initWorkspaceAndGetClient(config)
+	err := tokenPayload.initWorkspaceAndGetClient(config)
 	assert.NoError(t, err, err)
 	api := service.DBApiClient{}
 	api.SetConfig(config)
@@ -58,3 +54,10 @@ func TestAzureAuthCreateApiToken(t *testing.T) {
 
 	assert.NoError(t, instancePoolErr, instancePoolErr)
 }
+
+// getAndAssertEnv fetches the env for testing and also asserts that the env value is not Zero i.e ""
+func getAndAssertEnv(t *testing.T, key string) string {
+	value, present := os.LookupEnv(key)
+	assert.True(t, present, fmt.Sprintf("Env variable %s is not set", key))
+	return value
+}