Merge pull request #27 from databrickslabs/issue-21

Issue #21 Fix client_secret special characters

Author: stikkireddy

.gitignore

@@ -297,6 +297,8 @@ override.tf.json
 # http://iamzed.com/2009/05/07/a-primer-on-virtualenv/
 pyvenv.cfg
 .env
+.aws.env
+.azure.env
 .venv
 env/
 venv/

Makefile

@@ -32,7 +32,8 @@ build: lint test fmt
 
 lint:
 	@echo "==> Linting source code with golangci-lint..."
-	@golangci-lint run --skip-dirs-use-default --timeout 5m
+	@golangci-lint run --skip-dirs-use-default --timeout 5m --build-tags=azure
+	@golangci-lint run --skip-dirs-use-default --timeout 5m --build-tags=aws
 
 fmt: lint
 	@echo "==> Formatting source code with gofmt..."
@@ -61,10 +62,15 @@ vendor:
 	@echo "==> Filling vendor folder with library code..."
 	@go mod vendor
 
-# INTEGRATION TESTING WITH TERRAFORM EXAMPLES
-terraform-acc: fmt build
-	@echo "==> Running Terraform Acceptance Tests..."
-	@TF_ACC=1 gotestsum --format short-verbose --raw-command go test -v -json -short -coverprofile=coverage.out ./...
+# INTEGRATION TESTING WITH AZURE
+terraform-acc-azure: fmt
+	@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 ./...
+
+# INTEGRATION TESTING WITH AWS
+terraform-acc-aws: fmt
+	@echo "==> Running Terraform Acceptance Tests for AWS..."
+	@CLOUD_ENV="aws" TF_ACC=1 gotestsum --format short-verbose --raw-command go test -v -json -tags=aws  -short -coverprofile=coverage.out ./...
 
 terraform-setup: build
 	@echo "==> Initializing Terraform..."

README.md

@@ -147,6 +147,51 @@ $ docker run -it -v $(pwd):/workpace -w /workpace databricks-terraform plan
 $ docker run -it -v $(pwd):/workpace -w /workpace databricks-terraform apply
 ```
 
+## Testing
+
+* [ ] Integration tests should be run at a client level against both azure and aws to maintain sdk parity against both apis **(currently only on one cloud)**
+* [x] Terraform acceptance tests should be run against both aws and azure to maintain parity of provider between both cloud services **(currently only on one cloud)**
+
+### Linting
+
+Please use makefile for linting. If you run `golangci-lint` by itself it will fail due to different tags containing same functions. 
+So please run `make lint` instead.
+
+### Integration Testing
+
+Currently Databricks supports two cloud providers `azure` and `aws` thus integration testing with the correct cloud service provider is 
+crucial for making sure that the provider behaves as expected on all supported clouds. This type of testing separation is being managed via build tags 
+to allow for duplicate method names and environment variables to configure clients. 
+
+The current integration test implementation uses `CLOUD_ENV` environment variable and can use the value of `azure` or `aws`. 
+You can execute the acceptance with the following make commands `make terraform-acc-azure`, and `make terraform-acc-aws` for 
+azure and aws respectively. 
+
+This involves bootstrapping the provider via a .env configuration file. Without these files in the root directory the tests 
+will fail as the provider will not have a authorized token and host.
+
+The configuration file for `aws` should be like the following and be named `.aws.env`:
+```.env
+DATABRICKS_HOST=<host>
+DATABRICKS_TOKEN=<token>
+```
+
+The configuration file for `azure` should be like the following and be named `.azure.env`:
+```.env
+DATABRICKS_AZURE_CLIENT_ID=<enterprise app client id>
+DATABRICKS_AZURE_CLIENT_SECRET=<enterprise app client secret>
+DATABRICKS_AZURE_TENANT_ID=<azure ad tenant id>
+DATABRICKS_AZURE_SUBSCRIPTION_ID=<azure subscription id>
+DATABRICKS_AZURE_RESOURCE_GROUP=<resource group where the workspace is>
+AZURE_REGION=<region where the workspace is>
+DATABRICKS_AZURE_MANAGED_RESOURCE_GROUP=<azure databricks managed resource group for workspace>
+DATABRICKS_AZURE_WORKSPACE_NAME=<azure databricks workspace name>
+```
+
+Note that azure integration tests will use service principal based auth. Even though it is using a service principal, 
+it will still be generating a personal access token to perform creation of resources. 
+
+
 ## Project Components
 
 ### Databricks Terraform Provider Resources State
@@ -192,11 +237,6 @@ $ docker run -it -v $(pwd):/workpace -w /workpace databricks-terraform apply
 | databricks_table            | :white_large_square: | :white_large_square: | :white_large_square: | :white_large_square: |
 
 
-## Testing
-
-* [ ] Integration tests should be run at a client level against both azure and aws to maintain sdk parity against both apis **(currently only on one cloud)**
-* [ ] Terraform acceptance tests should be run against both aws and azure to maintain parity of provider between both cloud services **(currently only on one cloud)**
-
 ## Project Support
 Please note that all projects in the /databrickslabs github account are provided for your exploration only, and are not formally supported by Databricks with Service Level Agreements (SLAs).  They are provided AS-IS and we do not make any guarantees of any kind.  Please do not submit a support ticket relating to any issues arising from the use of these projects.
 

client/service/client.go

@@ -68,7 +68,7 @@ type DBApiClientConfig struct {
 // Setup initializes the client
 func (c *DBApiClientConfig) Setup() {
 	if c.TimeoutSeconds == 0 {
-		c.TimeoutSeconds = 10
+		c.TimeoutSeconds = 60
 	}
 	c.client = http.Client{
 		Timeout: time.Duration(time.Duration(c.TimeoutSeconds) * time.Second),

databricks/azure_ws_init.go renamed to databricks/azure_auth.go

@@ -2,15 +2,20 @@ package databricks
 
 import (
 	"encoding/json"
+	"fmt"
+	"github.com/Azure/go-autorest/autorest/adal"
+	"github.com/Azure/go-autorest/autorest/azure"
 	"github.com/databrickslabs/databricks-terraform/client/service"
 	"log"
 	"net/http"
+	urlParse "net/url"
+	"strings"
+	"time"
 )
 
 // List of management information
 const (
-	ManagementResourceEndpoint string = "https://management.core.windows.net/"
-	ADBResourceID              string = "2ff814a6-3304-4ab8-85cb-cd0e6f879c1d"
+	ADBResourceID string = "2ff814a6-3304-4ab8-85cb-cd0e6f879c1d"
 )
 
 // AzureAuth is a struct that contains information about the azure sp authentication
@@ -48,31 +53,36 @@ type WorkspaceRequest struct {
 
 func (a *AzureAuth) getManagementToken(config *service.DBApiClientConfig) error {
 	log.Println("[DEBUG] Creating Azure Databricks management OAuth token.")
-	url := "https://login.microsoftonline.com/" + a.TokenPayload.TenantID + "/oauth2/token"
-	payload := "grant_type=client_credentials&client_id=" + a.TokenPayload.ClientID + "&client_secret=" + a.TokenPayload.ClientSecret + "&resource=" + ManagementResourceEndpoint
-	headers := map[string]string{
-		"Content-Type":  "application/x-www-form-urlencoded",
-		"cache-control": "no-cache",
+	mgmtTokenOAuthCfg, err := adal.NewOAuthConfigWithAPIVersion(azure.PublicCloud.ActiveDirectoryEndpoint,
+		a.TokenPayload.TenantID,
+		nil)
+	if err != nil {
+		return err
 	}
-
-	var responseMap map[string]interface{}
-	resp, err := service.PerformQuery(config, http.MethodPost, url, "2.0", headers, false, true, payload, nil)
+	mgmtToken, err := adal.NewServicePrincipalToken(
+		*mgmtTokenOAuthCfg,
+		a.TokenPayload.ClientID,
+		a.TokenPayload.ClientSecret,
+		azure.PublicCloud.ServiceManagementEndpoint)
 	if err != nil {
 		return err
 	}
-	err = json.Unmarshal(resp, &responseMap)
+	err = mgmtToken.Refresh()
 	if err != nil {
 		return err
 	}
-	a.ManagementToken = responseMap["access_token"].(string)
+	a.ManagementToken = mgmtToken.OAuthToken()
 	return nil
 }
 
 func (a *AzureAuth) getWorkspaceID(config *service.DBApiClientConfig) error {
 	log.Println("[DEBUG] Getting Workspace ID via management token.")
-	url := "https://management.azure.com/subscriptions/" + a.TokenPayload.SubscriptionID + "/resourceGroups/" + a.TokenPayload.ResourceGroup + "/providers/Microsoft.Databricks/workspaces/" + a.TokenPayload.WorkspaceName + "" +
-		"?api-version=2018-04-01"
-
+	// Escape all the ids
+	url := fmt.Sprintf("https://management.azure.com/subscriptions/%s/resourceGroups/%s"+
+		"/providers/Microsoft.Databricks/workspaces/%s?api-version=2018-04-01",
+		urlParse.PathEscape(a.TokenPayload.SubscriptionID),
+		urlParse.PathEscape(a.TokenPayload.ResourceGroup),
+		urlParse.PathEscape(a.TokenPayload.WorkspaceName))
 	payload := &WorkspaceRequest{
 		Properties: &WsProps{ManagedResourceGroupID: "/subscriptions/" + a.TokenPayload.SubscriptionID + "/resourceGroups/" + a.TokenPayload.ManagedResourceGroup},
 		Name:       a.TokenPayload.WorkspaceName,
@@ -89,36 +99,37 @@ func (a *AzureAuth) getWorkspaceID(config *service.DBApiClientConfig) error {
 	if err != nil {
 		return err
 	}
+
 	err = json.Unmarshal(resp, &responseMap)
 	if err != nil {
 		return err
 	}
-	log.Println(responseMap)
 	a.AdbWorkspaceResourceID = responseMap["id"].(string)
 	return err
 }
 
 func (a *AzureAuth) getADBPlatformToken(clientConfig *service.DBApiClientConfig) error {
-	log.Println("[DEBUG] Creating Azure Databricks platform OAuth token.")
-	url := "https://login.microsoftonline.com/" + a.TokenPayload.TenantID + "/oauth2/token"
-	payload := "grant_type=client_credentials&client_id=" + a.TokenPayload.ClientID + "&client_secret=" + a.TokenPayload.ClientSecret + "&resource=" + ADBResourceID
-	headers := map[string]string{
-		"Content-Type":  "application/x-www-form-urlencoded",
-		"cache-control": "no-cache",
+	log.Println("[DEBUG] Creating Azure Databricks management OAuth token.")
+	platformTokenOAuthCfg, err := adal.NewOAuthConfigWithAPIVersion(azure.PublicCloud.ActiveDirectoryEndpoint,
+		a.TokenPayload.TenantID,
+		nil)
+	if err != nil {
+		return err
 	}
-
-	var responseMap map[string]interface{}
-	resp, err := service.PerformQuery(clientConfig, http.MethodPost, url, "2.0", headers, false, true, payload, &service.SecretsMask{
-		Secrets: []string{a.TokenPayload.ClientID, a.TokenPayload.TenantID, a.TokenPayload.ClientSecret},
-	})
+	platformToken, err := adal.NewServicePrincipalToken(
+		*platformTokenOAuthCfg,
+		a.TokenPayload.ClientID,
+		a.TokenPayload.ClientSecret,
+		ADBResourceID)
 	if err != nil {
 		return err
 	}
-	err = json.Unmarshal(resp, &responseMap)
+
+	err = platformToken.Refresh()
 	if err != nil {
 		return err
 	}
-	a.AdbPlatformToken = responseMap["access_token"].(string)
+	a.AdbPlatformToken = platformToken.OAuthToken()
 	return nil
 }
 
@@ -155,31 +166,100 @@ func (a *AzureAuth) getWorkspaceAccessToken(config *service.DBApiClientConfig) e
 	return 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
+// 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) (service.DBApiClient, error) {
 	var dbClient service.DBApiClient
+
+	// Get management token
 	err := a.getManagementToken(config)
 	if err != nil {
 		return dbClient, err
 	}
 
+	// Get workspace access token
 	err = a.getWorkspaceID(config)
 	if err != nil {
 		return dbClient, err
 	}
+
+	// Get platform token
 	err = a.getADBPlatformToken(config)
 	if err != nil {
 		return dbClient, err
 	}
+
+	// Get workspace personal access token
 	err = a.getWorkspaceAccessToken(config)
 	if err != nil {
 		return dbClient, err
 	}
 
 	var newOption service.DBApiClientConfig
-	newOption.Token = a.AdbAccessToken
+
+	// TODO: Eventually change this to include new Databricks domain names. May have to add new vars and/or deprecate existing args.
 	newOption.Host = "https://" + a.TokenPayload.AzureRegion + ".azuredatabricks.net"
+	newOption.Token = a.AdbAccessToken
+
+	// Headers to use aad tokens, hidden till tokens support secrets, scopes and acls
+	//newOption.DefaultHeaders = 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,
+	//	"cache-control":                            "no-cache",
+	//}
 	dbClient.SetConfig(&newOption)
 
-	_, err = dbClient.Clusters().ListNodeTypes()
+	// So when the workspace is initially created sometimes it fails to perform api calls so this is a simple test
+	// to verify that the workspace has been created successfully. The retry is intentional as sometimes after workspace
+	// creation the API's will not work correctly. This may also
+	err = validateWorkspaceApis(10, 30, func(attempt int) error {
+		_, err = dbClient.Clusters().ListNodeTypes()
+		return err
+	})
+
 	return dbClient, err
 }
+
+func validateWorkspaceApis(sleepDurationSeconds time.Duration, timeoutDurationMinutes time.Duration, do func(attempt int) error) error {
+	errChan := make(chan error, 1)
+	var timeoutBool = false
+	var attempts int
+	var expectedError error
+	apisAreNotYetReadErr := "com.databricks.backend.manager.util.UnknownWorkerEnvironmentException: Unknown worker environment WorkerEnvId"
+	go func(attempts *int, expectedError *error, timeout *bool) {
+		for {
+			err := do(*attempts)
+			// Timeout and terminate go routine so it does not leak
+			if *timeout {
+				errChan <- err
+				return
+			}
+			if err == nil {
+				errChan <- err
+				return
+			}
+			if !strings.Contains(err.Error(), apisAreNotYetReadErr) {
+				errChan <- err
+				return
+			}
+			log.Println(fmt.Sprintf("Waiting for cluster apis to not throw error: %v", err))
+			*attempts++
+			*expectedError = err
+			time.Sleep(sleepDurationSeconds * time.Second)
+		}
+	}(&attempts, &expectedError, &timeoutBool)
+	select {
+	case err := <-errChan:
+		if err == nil {
+			log.Printf("Returned nil error after %v attempts\n", attempts)
+		}
+		return err
+	case <-time.After(timeoutDurationMinutes * time.Minute):
+		timeoutBool = true
+		return fmt.Errorf("timed out waiting for ready state after %v attempts with error %v", attempts, expectedError)
+	}
+}