Verify workspace reachability through HTTP client

* `WorkspacesAPI.WaitForRunning` used to check if workspace is reachable through `net.Dial` on TCP level, though this scenario ignored HTTPS proxies, which are common for enterprise setup.
* `mws.Workspace` now gets `workspace_url` computed in accordance to Accounts API endpoint used from a single place in `WorkspacesAPI.Read`, if there's no `workspace_url` returned in the response, which allows fully unit testing the verification flow.
* added `DatabricksClient.ClientForHost` to talk to just created workspace with the same authentication parameters.

This fixes #874

Author: nfx

common/client.go

@@ -436,3 +436,29 @@ func (c *DatabricksClient) FormatURL(strs ...string) string {
 	data := append([]string{host}, strs...)
 	return strings.Join(data, "")
 }
+
+// ClientForHost creates a new DatabricksClient instance with the same auth parameters,
+// but for the given host. Authentication has to be reinitialized, as Google OIDC has
+// different authorizers, depending if it's workspace or Accounts API we're talking to.
+func (c *DatabricksClient) ClientForHost(url string) *DatabricksClient {
+	return &DatabricksClient{
+		Host:                 url,
+		Username:             c.Username,
+		Password:             c.Password,
+		Token:                c.Token,
+		Profile:              c.Profile,
+		ConfigFile:           c.ConfigFile,
+		GoogleServiceAccount: c.GoogleServiceAccount,
+		AzurermEnvironment:   c.AzurermEnvironment,
+		InsecureSkipVerify:   c.InsecureSkipVerify,
+		HTTPTimeoutSeconds:   c.HTTPTimeoutSeconds,
+		DebugTruncateBytes:   c.DebugTruncateBytes,
+		DebugHeaders:         c.DebugHeaders,
+		RateLimitPerSecond:   c.RateLimitPerSecond,
+		Provider:             c.Provider,
+		rateLimiter:          c.rateLimiter,
+		httpClient:           c.httpClient,
+		configAttributesUsed: c.configAttributesUsed,
+		commandFactory:       c.commandFactory,
+	}
+}

common/http.go

@@ -220,7 +220,11 @@ func (c *DatabricksClient) parseError(resp *http.Response) APIError {
 // checkHTTPRetry inspects HTTP errors from the Databricks API for known transient errors on Workspace creation
 func (c *DatabricksClient) checkHTTPRetry(ctx context.Context, resp *http.Response, err error) (bool, error) {
 	if ue, ok := err.(*url.Error); ok {
-		apiError := APIError{ErrorCode: "IO_ERROR", Message: ue.Error()}
+		apiError := APIError{
+			ErrorCode:  "IO_ERROR",
+			StatusCode: 523,
+			Message:    ue.Error(),
+		}
 		return apiError.IsRetriable(), apiError
 	}
 	if resp == nil {

mws/resource_workspace.go

@@ -52,16 +52,6 @@ func (a WorkspacesAPI) Create(ws *Workspace, timeout time.Duration) error {
 	return nil
 }
 
-func dial(hostAndPort, url string, timeout time.Duration) *resource.RetryError {
-	conn, err := net.DialTimeout("tcp", hostAndPort, timeout)
-	if err != nil {
-		return resource.RetryableError(err)
-	}
-	log.Printf("[INFO] Workspace %s is ready to use", url)
-	defer conn.Close()
-	return nil
-}
-
 // generateWorkspaceHostname computes the hostname for the specified workspace,
 // given the account console hostname.
 func generateWorkspaceHostname(client *common.DatabricksClient, ws Workspace) string {
@@ -85,6 +75,42 @@ func generateWorkspaceHostname(client *common.DatabricksClient, ws Workspace) st
 	return strings.Join(chunks, ".")
 }
 
+func (a WorkspacesAPI) verifyWorkspaceReachable(ws Workspace) *resource.RetryError {
+	ctx, cancel := context.WithTimeout(a.context, 10*time.Second)
+	defer cancel()
+	// wait for DNS caches to refresh, as sometimes we cannot make
+	// API calls to new workspaces immediately after it's created
+	wsClient := a.client.ClientForHost(ws.WorkspaceURL)
+	// make a request to Tokens API, just to verify there are no errors
+	var response map[string]interface{}
+	err := wsClient.Get(ctx, "/token/list", nil, &response)
+	if apiError, ok := err.(common.APIError); ok {
+		err = fmt.Errorf("workspace %s is not yet reachable: %s",
+			ws.WorkspaceURL, apiError)
+		log.Printf("[INFO] %s", err)
+		// expected to retry on: dial tcp: lookup XXX: no such host
+		return resource.RetryableError(err)
+	}
+	return nil
+}
+
+func (a WorkspacesAPI) explainWorkspaceFailure(ws Workspace) error {
+	if ws.NetworkID == "" {
+		return fmt.Errorf(ws.WorkspaceStatusMessage)
+	}
+	network, nerr := NewNetworksAPI(a.context, a.client).Read(ws.AccountID, ws.NetworkID)
+	if nerr != nil {
+		return fmt.Errorf("failed to start workspace. Cannot read network: %s", nerr)
+	}
+	var strBuffer bytes.Buffer
+	for _, networkHealth := range network.ErrorMessages {
+		strBuffer.WriteString(fmt.Sprintf("error: %s;error_msg: %s;",
+			networkHealth.ErrorType, networkHealth.ErrorMessage))
+	}
+	return fmt.Errorf("Workspace failed to create: %v, network error message: %v",
+		ws.WorkspaceStatusMessage, strBuffer.String())
+}
+
 // WaitForRunning will wait until workspace is running, otherwise will try to explain why it failed
 func (a WorkspacesAPI) WaitForRunning(ws Workspace, timeout time.Duration) error {
 	return resource.RetryContext(a.context, timeout, func() *resource.RetryError {
@@ -94,36 +120,17 @@ func (a WorkspacesAPI) WaitForRunning(ws Workspace, timeout time.Duration) error
 		}
 		switch workspace.WorkspaceStatus {
 		case WorkspaceStatusRunning:
-			// wait for DNS caches to refresh, as sometimes we cannot make
-			// API calls to new workspaces immediately after it's created
-			host := generateWorkspaceHostname(a.client, ws)
-			hostAndPort := fmt.Sprintf("%s:443", host)
-			url := fmt.Sprintf("https://%s", host)
 			log.Printf("[INFO] Workspace is now running")
-			if strings.Contains(workspace.DeploymentName, "900150983cd24fb0") {
+			if strings.Contains(ws.DeploymentName, "900150983cd24fb0") {
 				// nobody would probably name workspace as 900150983cd24fb0,
 				// so we'll use it as unit testing shim
 				return nil
 			}
-			return dial(hostAndPort, url, 10*time.Second)
+			return a.verifyWorkspaceReachable(workspace)
 		case WorkspaceStatusCanceled, WorkspaceStatusFailed:
 			log.Printf("[ERROR] Cannot start workspace: %s", workspace.WorkspaceStatusMessage)
-			if workspace.NetworkID == "" {
-				return resource.NonRetryableError(fmt.Errorf(workspace.WorkspaceStatusMessage))
-			}
-			network, nerr := NewNetworksAPI(a.context, a.client).Read(ws.AccountID, ws.NetworkID)
-			if nerr != nil {
-				return resource.NonRetryableError(fmt.Errorf(
-					"failed to start workspace. Cannot read network: %s", nerr))
-			}
-			var strBuffer bytes.Buffer
-			for _, networkHealth := range network.ErrorMessages {
-				strBuffer.WriteString(fmt.Sprintf("error: %s;error_msg: %s;",
-					networkHealth.ErrorType, networkHealth.ErrorMessage))
-			}
-			return resource.NonRetryableError(fmt.Errorf(
-				"Workspace failed to create: %v, network error message: %v",
-				workspace.WorkspaceStatusMessage, strBuffer.String()))
+			err = a.explainWorkspaceFailure(workspace)
+			return resource.NonRetryableError(err)
 		default:
 			log.Printf("[INFO] Workspace %s is %s: %s", workspace.DeploymentName,
 				workspace.WorkspaceStatus, workspace.WorkspaceStatusMessage)
@@ -162,6 +169,11 @@ func (a WorkspacesAPI) Read(mwsAcctID, workspaceID string) (Workspace, error) {
 	var mwsWorkspace Workspace
 	workspacesAPIPath := fmt.Sprintf("/accounts/%s/workspaces/%s", mwsAcctID, workspaceID)
 	err := a.client.Get(a.context, workspacesAPIPath, nil, &mwsWorkspace)
+	if err == nil && mwsWorkspace.WorkspaceURL == "" {
+		// generate workspace URL based on client's hostname, if response contains no URL
+		host := generateWorkspaceHostname(a.client, mwsWorkspace)
+		mwsWorkspace.WorkspaceURL = fmt.Sprintf("https://%s", host)
+	}
 	return mwsWorkspace, err
 }
 
@@ -296,7 +308,6 @@ func ResourceWorkspace() *schema.Resource {
 			// Default the value of `is_no_public_ip_enabled` because it isn't part of the GET payload.
 			// The field is only used on creation and we therefore suppress all diffs.
 			workspace.IsNoPublicIPEnabled = true
-			workspace.WorkspaceURL = fmt.Sprintf("https://%s", generateWorkspaceHostname(c, workspace))
 			if err = common.StructToData(workspace, workspaceSchema, d); err != nil {
 				return err
 			}

mws/resource_workspace_test.go

@@ -3,9 +3,6 @@ package mws
 import (
 	"context"
 	"fmt"
-	"net/http"
-	"net/http/httptest"
-	"strings"
 	"testing"
 	"time"
 
@@ -906,13 +903,60 @@ func TestListWorkspaces(t *testing.T) {
 	assert.Len(t, l, 0)
 }
 
-func TestDial(t *testing.T) {
-	err := dial("127.0.0.1:32456", "localhost", 50*time.Millisecond)
-	assert.NotNil(t, err)
-	assert.Equal(t, err.Err.Error(), "dial tcp 127.0.0.1:32456: connect: connection refused")
+func TestWorkspace_WaitForResolve_Failure(t *testing.T) {
+	qa.HTTPFixturesApply(t, []qa.HTTPFixture{
+		{
+			Method:       "GET",
+			Resource:     "/api/2.0/accounts/abc/workspaces/1234",
+			ReuseRequest: true,
+			Response: Workspace{
+				AccountID:       "abc",
+				WorkspaceID:     1234,
+				WorkspaceStatus: "RUNNING",
+				WorkspaceURL:    "https://foo-bar-baz.cloud.databricks.com",
+			},
+		},
+	}, func(ctx context.Context, client *common.DatabricksClient) {
+		a := NewWorkspacesAPI(ctx, client)
+		err := a.WaitForRunning(Workspace{
+			AccountID:   "abc",
+			WorkspaceID: 1234,
+		}, 1*time.Second)
+		assert.EqualError(t, err, "workspace https://foo-bar-baz.cloud.databricks.com is not yet reachable:"+
+			" Get \"https://foo-bar-baz.cloud.databricks.com/api/2.0/token/list\": "+
+			"dial tcp: lookup foo-bar-baz.cloud.databricks.com: no such host")
+	})
+}
 
-	s := httptest.NewServer(http.HandlerFunc(http.NotFound))
-	defer s.Close()
-	err = dial(strings.ReplaceAll(s.URL, "http://", ""), s.URL, 500*time.Millisecond)
-	assert.Nil(t, err)
+func TestWorkspace_WaitForResolve(t *testing.T) {
+	// outer HTTP server is used for inner request for "just created" workspace
+	qa.HTTPFixturesApply(t, []qa.HTTPFixture{
+		{
+			Method:   "GET",
+			Resource: "/api/2.0/token/list",
+			Response: `{}`, // we just need a JSON for this
+		},
+	}, func(ctx context.Context, wsClient *common.DatabricksClient) {
+		// inner HTTP server is used for outer request for Accounts API
+		qa.HTTPFixturesApply(t, []qa.HTTPFixture{
+			{
+				Method:       "GET",
+				Resource:     "/api/2.0/accounts/abc/workspaces/1234",
+				ReuseRequest: true,
+				Response: Workspace{
+					AccountID:       "abc",
+					WorkspaceID:     1234,
+					WorkspaceStatus: "RUNNING",
+					WorkspaceURL:    wsClient.Host,
+				},
+			},
+		}, func(ctx context.Context, client *common.DatabricksClient) {
+			a := NewWorkspacesAPI(ctx, client)
+			err := a.WaitForRunning(Workspace{
+				AccountID:   "abc",
+				WorkspaceID: 1234,
+			}, 1*time.Second)
+			assert.NoError(t, err)
+		})
+	})
 }