Merge pull request #27538 from DrFaust92/sagemaker-workforce-vpc-config

Sagemaker workforce - vpc support

.changelog/27538.txt

@@ -0,0 +1,3 @@
+```release-note:enhancement
+resource/aws_sagemaker_workforce: Add `workforce_vpc_config` argument
+```

internal/service/ec2/sweep.go

@@ -314,6 +314,7 @@ func init() {
 		F:    sweepVPCEndpoints,
 		Dependencies: []string{
 			"aws_route_table",
+			"aws_sagemaker_workforce",
 		},
 	})
 

internal/service/sagemaker/sagemaker_test.go

@@ -70,6 +70,7 @@ func TestAccSageMaker_serial(t *testing.T) {
 			"CognitoConfig":  testAccWorkforce_cognitoConfig,
 			"OidcConfig":     testAccWorkforce_oidcConfig,
 			"SourceIpConfig": testAccWorkforce_sourceIPConfig,
+			"VPC":            testAccWorkforce_vpc,
 		},
 		"Workteam": {
 			"disappears":         testAccWorkteam_disappears,

internal/service/sagemaker/status.go

@@ -248,3 +248,19 @@ func StatusProject(conn *sagemaker.SageMaker, name string) resource.StateRefresh
 		return output, aws.StringValue(output.ProjectStatus), nil
 	}
 }
+
+func StatusWorkforce(conn *sagemaker.SageMaker, name string) resource.StateRefreshFunc {
+	return func() (interface{}, string, error) {
+		output, err := FindWorkforceByName(conn, name)
+
+		if tfresource.NotFound(err) {
+			return nil, "", nil
+		}
+
+		if err != nil {
+			return nil, "", err
+		}
+
+		return output, aws.StringValue(output.Status), nil
+	}
+}

internal/service/sagemaker/wait.go

@@ -32,6 +32,8 @@ const (
 	FlowDefinitionDeletedTimeout      = 2 * time.Minute
 	ProjectCreatedTimeout             = 15 * time.Minute
 	ProjectDeletedTimeout             = 15 * time.Minute
+	WorkforceActiveTimeout            = 10 * time.Minute
+	WorkforceDeletedTimeout           = 10 * time.Minute
 )
 
 // WaitNotebookInstanceInService waits for a NotebookInstance to return InService
@@ -508,3 +510,45 @@ func WaitProjectUpdated(conn *sagemaker.SageMaker, name string) (*sagemaker.Desc
 
 	return nil, err
 }
+
+func WaitWorkforceActive(conn *sagemaker.SageMaker, name string) (*sagemaker.Workforce, error) {
+	stateConf := &resource.StateChangeConf{
+		Pending: []string{sagemaker.WorkforceStatusInitializing, sagemaker.WorkforceStatusUpdating},
+		Target:  []string{sagemaker.WorkforceStatusActive},
+		Refresh: StatusWorkforce(conn, name),
+		Timeout: WorkforceActiveTimeout,
+	}
+
+	outputRaw, err := stateConf.WaitForState()
+
+	if output, ok := outputRaw.(*sagemaker.Workforce); ok {
+		if status, reason := aws.StringValue(output.Status), aws.StringValue(output.FailureReason); status == sagemaker.WorkforceStatusFailed && reason != "" {
+			tfresource.SetLastError(err, errors.New(reason))
+		}
+
+		return output, err
+	}
+
+	return nil, err
+}
+
+func WaitWorkforceDeleted(conn *sagemaker.SageMaker, name string) (*sagemaker.Workforce, error) {
+	stateConf := &resource.StateChangeConf{
+		Pending: []string{sagemaker.WorkforceStatusDeleting},
+		Target:  []string{},
+		Refresh: StatusWorkforce(conn, name),
+		Timeout: WorkforceDeletedTimeout,
+	}
+
+	outputRaw, err := stateConf.WaitForState()
+
+	if output, ok := outputRaw.(*sagemaker.Workforce); ok {
+		if status, reason := aws.StringValue(output.Status), aws.StringValue(output.FailureReason); status == sagemaker.WorkforceStatusFailed && reason != "" {
+			tfresource.SetLastError(err, errors.New(reason))
+		}
+
+		return output, err
+	}
+
+	return nil, err
+}

internal/service/sagemaker/workforce.go

@@ -146,6 +146,35 @@ func ResourceWorkforce() *schema.Resource {
 					validation.StringMatch(regexp.MustCompile(`^[a-zA-Z0-9]([a-zA-Z0-9\-])*$`), "Valid characters are a-z, A-Z, 0-9, and - (hyphen)."),
 				),
 			},
+			"workforce_vpc_config": {
+				Type:     schema.TypeList,
+				Optional: true,
+				MaxItems: 1,
+				Elem: &schema.Resource{
+					Schema: map[string]*schema.Schema{
+						"security_group_ids": {
+							Type:     schema.TypeSet,
+							Optional: true,
+							MaxItems: 5,
+							Elem:     &schema.Schema{Type: schema.TypeString},
+						},
+						"subnets": {
+							Type:     schema.TypeSet,
+							Optional: true,
+							MaxItems: 16,
+							Elem:     &schema.Schema{Type: schema.TypeString},
+						},
+						"vpc_endpoint_id": {
+							Type:     schema.TypeString,
+							Computed: true,
+						},
+						"vpc_id": {
+							Type:     schema.TypeString,
+							Optional: true,
+						},
+					},
+				},
+			},
 		},
 	}
 }
@@ -170,7 +199,10 @@ func resourceWorkforceCreate(d *schema.ResourceData, meta interface{}) error {
 		input.SourceIpConfig = expandWorkforceSourceIPConfig(v.([]interface{}))
 	}
 
-	log.Printf("[DEBUG] Creating SageMaker Workforce: %s", input)
+	if v, ok := d.GetOk("workforce_vpc_config"); ok {
+		input.WorkforceVpcConfig = expandWorkforceVPCConfig(v.([]interface{}))
+	}
+
 	_, err := conn.CreateWorkforce(input)
 
 	if err != nil {
@@ -179,6 +211,10 @@ func resourceWorkforceCreate(d *schema.ResourceData, meta interface{}) error {
 
 	d.SetId(name)
 
+	if _, err := WaitWorkforceActive(conn, name); err != nil {
+		return fmt.Errorf("waiting for SageMaker Workforce (%s) create: %w", d.Id(), err)
+	}
+
 	return resourceWorkforceRead(d, meta)
 }
 
@@ -215,6 +251,10 @@ func resourceWorkforceRead(d *schema.ResourceData, meta interface{}) error {
 		return fmt.Errorf("setting source_ip_config: %w", err)
 	}
 
+	if err := d.Set("workforce_vpc_config", flattenWorkforceVPCConfig(workforce.WorkforceVpcConfig)); err != nil {
+		return fmt.Errorf("setting workforce_vpc_config: %w", err)
+	}
+
 	return nil
 }
 
@@ -233,13 +273,20 @@ func resourceWorkforceUpdate(d *schema.ResourceData, meta interface{}) error {
 		input.OidcConfig = expandWorkforceOIDCConfig(d.Get("oidc_config").([]interface{}))
 	}
 
-	log.Printf("[DEBUG] Updating SageMaker Workforce: %s", input)
+	if d.HasChange("workforce_vpc_config") {
+		input.WorkforceVpcConfig = expandWorkforceVPCConfig(d.Get("workforce_vpc_config").([]interface{}))
+	}
+
 	_, err := conn.UpdateWorkforce(input)
 
 	if err != nil {
 		return fmt.Errorf("updating SageMaker Workforce (%s): %w", d.Id(), err)
 	}
 
+	if _, err := WaitWorkforceActive(conn, d.Id()); err != nil {
+		return fmt.Errorf("waiting for SageMaker Workforce (%s) update: %w", d.Id(), err)
+	}
+
 	return resourceWorkforceRead(d, meta)
 }
 
@@ -259,6 +306,10 @@ func resourceWorkforceDelete(d *schema.ResourceData, meta interface{}) error {
 		return fmt.Errorf("deleting SageMaker Workforce (%s): %w", d.Id(), err)
 	}
 
+	if _, err := WaitWorkforceDeleted(conn, d.Id()); err != nil {
+		return fmt.Errorf("waiting for SageMaker Workforce (%s) delete: %w", d.Id(), err)
+	}
+
 	return nil
 }
 
@@ -355,3 +406,34 @@ func flattenWorkforceOIDCConfig(config *sagemaker.OidcConfigForResponse, clientS
 
 	return []map[string]interface{}{m}
 }
+
+func expandWorkforceVPCConfig(l []interface{}) *sagemaker.WorkforceVpcConfigRequest {
+	if len(l) == 0 || l[0] == nil {
+		return &sagemaker.WorkforceVpcConfigRequest{}
+	}
+
+	m := l[0].(map[string]interface{})
+
+	config := &sagemaker.WorkforceVpcConfigRequest{
+		SecurityGroupIds: flex.ExpandStringSet(m["security_group_ids"].(*schema.Set)),
+		Subnets:          flex.ExpandStringSet(m["subnets"].(*schema.Set)),
+		VpcId:            aws.String(m["vpc_id"].(string)),
+	}
+
+	return config
+}
+
+func flattenWorkforceVPCConfig(config *sagemaker.WorkforceVpcConfigResponse) []map[string]interface{} {
+	if config == nil {
+		return []map[string]interface{}{}
+	}
+
+	m := map[string]interface{}{
+		"security_group_ids": flex.FlattenStringSet(config.SecurityGroupIds),
+		"subnets":            flex.FlattenStringSet(config.Subnets),
+		"vpc_endpoint_id":    aws.StringValue(config.VpcEndpointId),
+		"vpc_id":             aws.StringValue(config.VpcId),
+	}
+
+	return []map[string]interface{}{m}
+}