Add expected_workspace_status in databricks_mws_workspaces to support least-privileged workspaces

NEXT_CHANGELOG.md

@@ -6,6 +6,8 @@
 
 ### New Features and Improvements
 
+* Added `expected_workspace_status` to `databricks_mws_workspaces` to support creating workspaces in provisioning status ([#5019](https://github.com/databricks/terraform-provider-databricks/pull/5019))
+
 ### Bug Fixes
 
 ### Documentation

docs/resources/mws_workspaces.md

@@ -309,6 +309,7 @@ The following arguments are available:
 * `custom_tags` - (Optional / AWS only) - The custom tags key-value pairing that is attached to this workspace. These tags will be applied to clusters automatically in addition to any `default_tags` or `custom_tags` on a cluster level. Please note it can take up to an hour for custom_tags to be set due to scheduling on Control Plane. After custom tags are applied, they can be modified however they can never be completely removed.
 * `pricing_tier` - (Optional) - The pricing tier of the workspace.
 * `compute_mode` - (Optional) - The compute mode for the workspace. When unset, a classic workspace is created, and both `credentials_id` and `storage_configuration_id` must be specified. When set to `SERVERLESS`, the resulting workspace is a serverless workspace, and `credentials_id` and `storage_configuration_id` must not be set. The only allowed value for this is `SERVERLESS`. Changing this field requires recreation of the workspace.
+* `expected_workspace_status` - (Optional / GCP only / Private Preview) - The expected status of the workspace. When unset, it defaults to `RUNNING`. When set to `PROVISIONING`, workspace provisioning will pause and not enter `RUNNING` status. The only allowed values for this is `RUNNING` and `PROVISIONING`.
 
 ~> Databricks strongly recommends using OAuth instead of PATs for user account client authentication and authorization due to the improved security
 

mws/mws_workspaces_test.go

@@ -258,6 +258,153 @@ func TestMwsAccGcpWorkspaces(t *testing.T) {
 	})
 }
 
+func TestMwsAccGcpWorkspacesProvisioningToRunning(t *testing.T) {
+	acceptance.AccountLevel(t, acceptance.Step{
+		Template: `
+		resource "databricks_mws_workspaces" "this" {
+			account_id                = "{env.DATABRICKS_ACCOUNT_ID}"
+			workspace_name            = "{env.TEST_PREFIX}-{var.STICKY_RANDOM}"
+			location                  = "{env.GOOGLE_REGION}"
+			expected_workspace_status = "PROVISIONING"
+
+			cloud_resource_container {
+				gcp {
+					project_id = "{env.GOOGLE_PROJECT}"
+				}
+			}
+		}`,
+		Check: func(s *terraform.State) error {
+			rs, ok := s.RootModule().Resources["databricks_mws_workspaces.this"]
+			if !ok {
+				return fmt.Errorf("databricks_mws_workspaces.this not found")
+			}
+			if rs.Primary.Attributes["workspace_id"] == "" {
+				return fmt.Errorf("workspace_id is empty")
+			}
+
+			expectedStatus := "PROVISIONING"
+			if status := rs.Primary.Attributes["workspace_status"]; status != expectedStatus {
+				return fmt.Errorf("expected workspace_status to be %s, got %s", expectedStatus, status)
+			}
+			return nil
+		},
+	}, acceptance.Step{
+		Template: `
+		resource "databricks_mws_networks" "this" {
+			account_id   = "{env.DATABRICKS_ACCOUNT_ID}"
+			network_name = "network-{var.STICKY_RANDOM}"
+			gcp_network_info {
+				network_project_id = "{env.GOOGLE_PROJECT}"
+				vpc_id = "{env.TEST_VPC_ID}"
+				subnet_id = "{env.TEST_SUBNET_ID}"
+				subnet_region = "{env.GOOGLE_REGION}"
+			}
+		}
+		resource "databricks_mws_workspaces" "this" {
+			account_id                = "{env.DATABRICKS_ACCOUNT_ID}"
+			workspace_name            = "{env.TEST_PREFIX}-{var.STICKY_RANDOM}"
+			location                  = "{env.GOOGLE_REGION}"
+			expected_workspace_status = "RUNNING"
+
+			cloud_resource_container {
+				gcp {
+					project_id = "{env.GOOGLE_PROJECT}"
+				}
+			}
+
+			network_id = databricks_mws_networks.this.network_id
+		}`,
+		Check: func(s *terraform.State) error {
+			rs, ok := s.RootModule().Resources["databricks_mws_workspaces.this"]
+			if !ok {
+				return fmt.Errorf("databricks_mws_workspaces.this not found")
+			}
+			if rs.Primary.Attributes["workspace_id"] == "" {
+				return fmt.Errorf("workspace_id is empty")
+			}
+
+			expectedStatus := "RUNNING"
+			if status := rs.Primary.Attributes["workspace_status"]; status != expectedStatus {
+				return fmt.Errorf("expected workspace_status to be %s, got %s", expectedStatus, status)
+			}
+			return nil
+		},
+	})
+}
+
+func TestMwsAccGcpWorkspacesUnsetExpectedState(t *testing.T) {
+	acceptance.AccountLevel(t, acceptance.Step{
+		Template: `
+		resource "databricks_mws_workspaces" "this" {
+			account_id                = "{env.DATABRICKS_ACCOUNT_ID}"
+			workspace_name            = "{env.TEST_PREFIX}-{var.STICKY_RANDOM}"
+			location                  = "{env.GOOGLE_REGION}"
+			expected_workspace_status = "PROVISIONING"
+
+			cloud_resource_container {
+				gcp {
+					project_id = "{env.GOOGLE_PROJECT}"
+				}
+			}
+		}`,
+		Check: func(s *terraform.State) error {
+			rs, ok := s.RootModule().Resources["databricks_mws_workspaces.this"]
+			if !ok {
+				return fmt.Errorf("databricks_mws_workspaces.this not found")
+			}
+			if rs.Primary.Attributes["workspace_id"] == "" {
+				return fmt.Errorf("workspace_id is empty")
+			}
+
+			expectedStatus := "PROVISIONING"
+			if status := rs.Primary.Attributes["workspace_status"]; status != expectedStatus {
+				return fmt.Errorf("expected workspace_status to be %s, got %s", expectedStatus, status)
+			}
+			return nil
+		},
+	}, acceptance.Step{
+		Template: `
+		resource "databricks_mws_networks" "this" {
+			account_id   = "{env.DATABRICKS_ACCOUNT_ID}"
+			network_name = "network-{var.STICKY_RANDOM}"
+			gcp_network_info {
+				network_project_id = "{env.GOOGLE_PROJECT}"
+				vpc_id = "{env.TEST_VPC_ID}"
+				subnet_id = "{env.TEST_SUBNET_ID}"
+				subnet_region = "{env.GOOGLE_REGION}"
+			}
+		}
+		resource "databricks_mws_workspaces" "this" {
+			account_id      = "{env.DATABRICKS_ACCOUNT_ID}"
+			workspace_name  = "{env.TEST_PREFIX}-{var.STICKY_RANDOM}"
+			location        = "{env.GOOGLE_REGION}"
+
+			cloud_resource_container {
+				gcp {
+					project_id = "{env.GOOGLE_PROJECT}"
+				}
+			}
+
+			network_id = databricks_mws_networks.this.network_id
+		}`,
+		Check: func(s *terraform.State) error {
+			rs, ok := s.RootModule().Resources["databricks_mws_workspaces.this"]
+			if !ok {
+				return fmt.Errorf("databricks_mws_workspaces.this not found")
+			}
+			if rs.Primary.Attributes["workspace_id"] == "" {
+				return fmt.Errorf("workspace_id is empty")
+			}
+
+			expectedStatus := "RUNNING"
+			if status := rs.Primary.Attributes["workspace_status"]; status != expectedStatus {
+				return fmt.Errorf("expected workspace_status to be %s, got %s", expectedStatus, status)
+			}
+			return nil
+		},
+	})
+}
+
 func TestMwsAccGcpByovpcWorkspaces(t *testing.T) {
 	acceptance.AccountLevel(t, acceptance.Step{
 		Template: `

mws/resource_mws_workspaces.go

@@ -97,6 +97,7 @@ type Workspace struct {
 	WorkspaceURL                        string                   `json:"workspace_url,omitempty" tf:"computed"`
 	WorkspaceStatus                     string                   `json:"workspace_status,omitempty" tf:"computed"`
 	WorkspaceStatusMessage              string                   `json:"workspace_status_message,omitempty" tf:"computed"`
+	ExpectedWorkspaceStatus             string                   `json:"expected_workspace_status,omitempty"`
 	CreationTime                        int64                    `json:"creation_time,omitempty" tf:"computed"`
 	ExternalCustomerInfo                *externalCustomerInfo    `json:"external_customer_info,omitempty"`
 	CloudResourceBucket                 *CloudResourceContainer  `json:"cloud_resource_container,omitempty"`
@@ -147,6 +148,9 @@ func (w *Workspace) MarshalJSON() ([]byte, error) {
 	if w.ComputeMode != "" {
 		workspaceCreationRequest["compute_mode"] = w.ComputeMode
 	}
+	if w.ExpectedWorkspaceStatus != "" {
+		workspaceCreationRequest["expected_workspace_status"] = w.ExpectedWorkspaceStatus
+	}
 	return json.Marshal(workspaceCreationRequest)
 }
 
@@ -161,7 +165,7 @@ func (a WorkspacesAPI) Create(ws *Workspace, timeout time.Duration) error {
 	if err != nil {
 		return err
 	}
-	if err = a.WaitForRunning(*ws, timeout); err != nil {
+	if err = a.WaitForExpectedStatus(*ws, ws.ExpectedWorkspaceStatus, timeout); err != nil {
 		log.Printf("[ERROR] Deleting failed workspace: %s", err)
 		if derr := a.Delete(ws.AccountID, fmt.Sprintf("%d", ws.WorkspaceID)); derr != nil {
 			return fmt.Errorf("%s - %s", err, derr)
@@ -239,22 +243,35 @@ func (a WorkspacesAPI) explainWorkspaceFailure(ws Workspace) error {
 		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 {
+// If expected_workspace_status is specified, WaitForExpectedStatus will wait until workspace is in the expected status.
+// If not, it will wait until workspace is running, and otherwise will try to explain why it failed.
+func (a WorkspacesAPI) WaitForExpectedStatus(ws Workspace, expectedStatus string, timeout time.Duration) error {
+	// If expected_status is empty, default to RUNNING
+	if expectedStatus == "" {
+		expectedStatus = WorkspaceStatusRunning
+		log.Printf("[INFO] No expected_workspace_status specified, defaulting to %s", expectedStatus)
+	}
+
 	return resource.RetryContext(a.context, timeout, func() *resource.RetryError {
+
 		workspace, err := a.Read(ws.AccountID, fmt.Sprintf("%d", ws.WorkspaceID))
 		if err != nil {
 			return resource.NonRetryableError(err)
 		}
+
 		switch workspace.WorkspaceStatus {
-		case WorkspaceStatusRunning:
-			log.Printf("[INFO] Workspace is now running")
-			if strings.Contains(ws.DeploymentName, "900150983cd24fb0") {
-				// nobody would probably name workspace as 900150983cd24fb0,
-				// so we'll use it as unit testing shim
-				return nil
+		case expectedStatus:
+			log.Printf("[INFO] Workspace is now in expected status %s", expectedStatus)
+			// only verify that workspace is reachable if expected status is RUNNING
+			if expectedStatus == WorkspaceStatusRunning {
+				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 a.verifyWorkspaceReachable(workspace)
 			}
-			return a.verifyWorkspaceReachable(workspace)
+			return nil
 		case WorkspaceStatusCanceled, WorkspaceStatusFailed:
 			log.Printf("[ERROR] Cannot start workspace: %s", workspace.WorkspaceStatusMessage)
 			err = a.explainWorkspaceFailure(workspace)
@@ -267,7 +284,15 @@ func (a WorkspacesAPI) WaitForRunning(ws Workspace, timeout time.Duration) error
 	})
 }
 
-var workspaceRunningUpdatesAllowed = []string{"credentials_id", "network_id", "storage_customer_managed_key_id", "private_access_settings_id", "managed_services_customer_managed_key_id", "custom_tags"}
+var workspaceRunningUpdatesAllowed = []string{
+	"credentials_id",
+	"network_id",
+	"storage_customer_managed_key_id",
+	"private_access_settings_id",
+	"managed_services_customer_managed_key_id",
+	"custom_tags",
+	"expected_workspace_status",
+}
 
 // UpdateRunning will update running workspace with couple of possible fields
 func (a WorkspacesAPI) UpdateRunning(ws Workspace, timeout time.Duration) error {
@@ -301,6 +326,9 @@ func (a WorkspacesAPI) UpdateRunning(ws Workspace, timeout time.Duration) error
 		}
 		request["custom_tags"] = ws.CustomTags
 	}
+	if ws.ExpectedWorkspaceStatus != "" {
+		request["expected_workspace_status"] = ws.ExpectedWorkspaceStatus
+	}
 
 	if len(request) == 0 {
 		return nil
@@ -310,7 +338,7 @@ func (a WorkspacesAPI) UpdateRunning(ws Workspace, timeout time.Duration) error
 	if err != nil {
 		return err
 	}
-	return a.WaitForRunning(ws, timeout)
+	return a.WaitForExpectedStatus(ws, ws.ExpectedWorkspaceStatus, timeout)
 }
 
 // Read will return the mws workspace metadata and status of the workspace deployment
@@ -551,6 +579,15 @@ func ResourceMwsWorkspaces() common.Resource {
 				Type:     schema.TypeString,
 				Computed: true,
 			}
+			// validate that expected_workspace_status is one of [PROVISIONING, RUNNING]
+			if f, ok := s["expected_workspace_status"]; ok {
+				f.ValidateDiagFunc = validation.ToDiagFunc(
+					validation.StringInSlice([]string{
+						WorkspaceStatusProvisioning,
+						WorkspaceStatusRunning,
+					}, false),
+				)
+			}
 			docOptions := docs.DocOptions{
 				Section:  docs.Guides,
 				Slug:     "gcp-workspace",
@@ -638,7 +675,13 @@ func ResourceMwsWorkspaces() common.Resource {
 			if err = common.StructToData(workspace, workspaceSchema, d); err != nil {
 				return err
 			}
-			err = workspacesAPI.WaitForRunning(workspace, d.Timeout(schema.TimeoutRead))
+			// The expected_workspace_status field is input only.
+			// Therefore, we need to read it from the original Terraform configuration.
+			expectedStatus := d.Get("expected_workspace_status").(string)
+			// PROVISIONING workspace import may fail because the "expected_workspace_status" is not included in the state during import, nor is it returned by the API.
+			// As a result, the provider will wait for RUNNING state, which will never happen, and timeout.
+			// TODO: fix this.
+			err = workspacesAPI.WaitForExpectedStatus(workspace, expectedStatus, d.Timeout(schema.TimeoutRead))
 			if err != nil {
 				return err
 			}