Allow custom tags on AWS customer managed VPC workspaces

docs/resources/mws_workspaces.md

@@ -192,6 +192,13 @@ resource "databricks_mws_workspaces" "this" {
   storage_configuration_id = databricks_mws_storage_configurations.this.storage_configuration_id
 
   token {}
+
+  # Optional Custom Tags
+  custom_tags = {
+
+    "SoldToCode" = "1234"
+
+  }
 }
 
 output "databricks_token" {
@@ -328,6 +335,7 @@ The following arguments are available:
   * `connectivity_type`: Specifies the network connectivity types for the GKE nodes and the GKE master network. Possible values are: `PRIVATE_NODE_PUBLIC_MASTER`, `PUBLIC_NODE_PUBLIC_MASTER`.
   * `master_ip_range`: The IP range from which to allocate GKE cluster master resources. This field will be ignored if GKE private cluster is not enabled. It must be exactly as big as `/28`.
 * `private_access_settings_id` - (Optional) Canonical unique identifier of [databricks_mws_private_access_settings](mws_private_access_settings.md) in Databricks Account.
+* `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.
 
 ### token block
 
@@ -346,6 +354,7 @@ On AWS, the following arguments could be modified after the workspace is running
 * `credentials_id`
 * `storage_customer_managed_key_id`
 * `private_access_settings_id`
+* `custom_tags`
 
 ## Attribute Reference
 
@@ -357,6 +366,7 @@ In addition to all arguments above, the following attributes are exported:
 * `workspace_status` - (String) workspace status
 * `creation_time` - (Integer) time when workspace was created
 * `workspace_url` - (String) URL of the workspace
+* `custom_tags` - (Map) Custom Tags (if present) added to workspace
 
 ## Import
 

internal/acceptance/mws_workspaces_test.go

@@ -53,6 +53,10 @@ func TestMwsAccWorkspaces(t *testing.T) {
 			storage_configuration_id = databricks_mws_storage_configurations.this.storage_configuration_id
 			managed_services_customer_managed_key_id = databricks_mws_customer_managed_keys.this.customer_managed_key_id
 
+			custom_tags = {
+				"randomkey" = "randomvalue"
+			}
+
 			token {
 				comment = "Test {var.RANDOM}"
 			}

mws/resource_mws_workspaces.go

@@ -100,6 +100,7 @@ type Workspace struct {
 	GkeConfig                           *GkeConfig               `json:"gke_config,omitempty" tf:"suppress_diff"`
 	Cloud                               string                   `json:"cloud,omitempty" tf:"computed"`
 	Location                            string                   `json:"location,omitempty"`
+	CustomTags                          map[string]string        `json:"custom_tags,omitempty"` // Optional for AWS, not allowed for GCP
 }
 
 // this type alias hack is required for Marshaller to work without an infinite loop
@@ -257,12 +258,12 @@ 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"}
+var workspaceRunningUpdatesAllowed = []string{"credentials_id", "network_id", "storage_customer_managed_key_id", "private_access_settings_id", "managed_services_customer_managed_key_id", "custom_tags"}
 
 // UpdateRunning will update running workspace with couple of possible fields
 func (a WorkspacesAPI) UpdateRunning(ws Workspace, timeout time.Duration) error {
 	workspacesAPIPath := fmt.Sprintf("/accounts/%s/workspaces/%d", ws.AccountID, ws.WorkspaceID)
-	request := map[string]string{}
+	request := map[string]any{}
 
 	if ws.CredentialsID != "" {
 		request["credentials_id"] = ws.CredentialsID
@@ -282,6 +283,12 @@ func (a WorkspacesAPI) UpdateRunning(ws Workspace, timeout time.Duration) error
 	if ws.StorageCustomerManagedKeyID != "" {
 		request["storage_customer_managed_key_id"] = ws.StorageCustomerManagedKeyID
 	}
+	if ws.CustomTags != nil {
+		if !a.client.IsAws() {
+			return fmt.Errorf("custom_tags are only allowed for AWS workspaces")
+		}
+		request["custom_tags"] = ws.CustomTags
+	}
 
 	if len(request) == 0 {
 		return nil
@@ -532,6 +539,9 @@ func ResourceMwsWorkspaces() *schema.Resource {
 			if err := requireFields(c.IsGcp(), d, "location"); err != nil {
 				return err
 			}
+			if !c.IsAws() && workspace.CustomTags != nil {
+				return fmt.Errorf("custom_tags are only allowed for AWS workspaces")
+			}
 			if len(workspace.CustomerManagedKeyID) > 0 && len(workspace.ManagedServicesCustomerManagedKeyID) == 0 {
 				log.Print("[INFO] Using existing customer_managed_key_id as value for new managed_services_customer_managed_key_id")
 				workspace.ManagedServicesCustomerManagedKeyID = workspace.CustomerManagedKeyID
@@ -859,6 +869,10 @@ func workspaceSchemaV2() cty.Type {
 					},
 				},
 			},
+			"custom_tags": {
+				Type:     schema.TypeMap,
+				Optional: true,
+			},
 		},
 	}).CoreConfigSchema().ImpliedType()
 }

mws/resource_mws_workspaces_test.go

@@ -34,6 +34,9 @@ func TestResourceWorkspaceCreate(t *testing.T) {
 					NetworkID:                           "fgh",
 					ManagedServicesCustomerManagedKeyID: "def",
 					StorageCustomerManagedKeyID:         "def",
+					CustomTags: map[string]string{
+						"SoldToCode": "1234",
+					},
 				},
 				Response: Workspace{
 					WorkspaceID:    1234,
@@ -57,6 +60,9 @@ func TestResourceWorkspaceCreate(t *testing.T) {
 					ManagedServicesCustomerManagedKeyID: "def",
 					StorageCustomerManagedKeyID:         "def",
 					AccountID:                           "abc",
+					CustomTags: map[string]string{
+						"SoldToCode": "1234",
+					},
 				},
 			},
 		},
@@ -71,6 +77,9 @@ func TestResourceWorkspaceCreate(t *testing.T) {
 			"workspace_name":                           "labdata",
 			"network_id":                               "fgh",
 			"storage_configuration_id":                 "ghi",
+			"custom_tags": map[string]any{
+				"SoldToCode": "1234",
+			},
 		},
 		Create: true,
 	}.Apply(t)
@@ -153,6 +162,87 @@ func TestResourceWorkspaceCreateGcp(t *testing.T) {
 	}.ApplyNoError(t)
 }
 
+func TestResourceWorkspaceCreate_Error_Custom_tags(t *testing.T) {
+	qa.ResourceFixture{
+		Fixtures: []qa.HTTPFixture{
+			{
+				Method:   "POST",
+				Resource: "/api/2.0/accounts/abc/workspaces",
+				// retreating to raw JSON, as certain fields don't work well together
+				ExpectedRequest: map[string]any{
+					"account_id": "abc",
+					"cloud":      "gcp",
+					"cloud_resource_container": map[string]any{
+						"gcp": map[string]any{
+							"project_id": "def",
+						},
+					},
+					"location":                   "bcd",
+					"private_access_settings_id": "pas_id_a",
+					"network_id":                 "net_id_a",
+					"gke_config": map[string]any{
+						"master_ip_range":   "e",
+						"connectivity_type": "PRIVATE_NODE_PUBLIC_MASTER",
+					},
+					"gcp_managed_network_config": map[string]any{
+						"gke_cluster_pod_ip_range":     "b",
+						"gke_cluster_service_ip_range": "c",
+						"subnet_cidr":                  "a",
+					},
+					"workspace_name": "labdata",
+					"custom_tags": map[string]any{
+						"SoldToCode": "1234",
+					},
+				},
+				Response: apierr.APIErrorBody{
+					ErrorCode: "INVALID_PARAMETER_VALUE",
+					Message:   "custom_tags are only allowed for AWS workspaces",
+				},
+			},
+			{
+				Method:       "GET",
+				ReuseRequest: true,
+				Resource:     "/api/2.0/accounts/abc/workspaces/1234",
+				Response: Workspace{
+					AccountID:       "abc",
+					WorkspaceID:     1234,
+					WorkspaceStatus: WorkspaceStatusRunning,
+					DeploymentName:  "900150983cd24fb0",
+					WorkspaceName:   "labdata",
+				},
+			},
+		},
+		Resource: ResourceMwsWorkspaces(),
+		HCL: `
+		account_id      = "abc"
+		workspace_name  = "labdata"
+		deployment_name = "900150983cd24fb0"
+		location        = "bcd"
+		cloud_resource_container {
+			gcp {
+				project_id = "def"
+			}
+		}
+		private_access_settings_id = "pas_id_a"
+		network_id = "net_id_a"
+		gcp_managed_network_config {
+			subnet_cidr = "a"
+			gke_cluster_pod_ip_range = "b"
+			gke_cluster_service_ip_range = "c"
+		}
+		gke_config {
+			connectivity_type = "PRIVATE_NODE_PUBLIC_MASTER"
+			master_ip_range = "e"
+		}
+		custom_tags = {
+			SoldToCode = "1234"
+		}
+		`,
+		Gcp:    true,
+		Create: true,
+	}.ExpectError(t, "custom_tags are only allowed for AWS workspaces")
+}
+
 func TestResourceWorkspaceCreateGcpPsc(t *testing.T) {
 	qa.ResourceFixture{
 		Fixtures: []qa.HTTPFixture{