feature(sagemaker notebook): Add DirectInternetAccess support

aws/resource_aws_sagemaker_notebook_instance.go

@@ -9,6 +9,7 @@ import (
 	"github.com/aws/aws-sdk-go/service/sagemaker"
 	"github.com/hashicorp/terraform-plugin-sdk/helper/resource"
 	"github.com/hashicorp/terraform-plugin-sdk/helper/schema"
+	"github.com/hashicorp/terraform-plugin-sdk/helper/validation"
 )
 
 func resourceAwsSagemakerNotebookInstance() *schema.Resource {
@@ -72,6 +73,16 @@ func resourceAwsSagemakerNotebookInstance() *schema.Resource {
 				ForceNew: true,
 			},
 
+			"direct_internet_access": {
+				Type:     schema.TypeString,
+				Optional: true,
+				ForceNew: true,
+				ValidateFunc: validation.StringInSlice([]string{
+					sagemaker.DirectInternetAccessDisabled,
+					sagemaker.DirectInternetAccessEnabled,
+				}, false),
+			},
+
 			"tags": tagsSchema(),
 		},
 	}
@@ -91,6 +102,9 @@ func resourceAwsSagemakerNotebookInstanceCreate(d *schema.ResourceData, meta int
 
 	if s, ok := d.GetOk("subnet_id"); ok {
 		createOpts.SubnetId = aws.String(s.(string))
+		if dia, ok := d.GetOk("direct_internet_access"); ok {
+			createOpts.DirectInternetAccess = aws.String(dia.(string))
+		}
 	}
 
 	if k, ok := d.GetOk("kms_key_id"); ok {
@@ -177,6 +191,11 @@ func resourceAwsSagemakerNotebookInstanceRead(d *schema.ResourceData, meta inter
 	if err := d.Set("arn", notebookInstance.NotebookInstanceArn); err != nil {
 		return fmt.Errorf("error setting arn for sagemaker notebook instance (%s): %s", d.Id(), err)
 	}
+
+	if err := d.Set("direct_internet_access", notebookInstance.DirectInternetAccess); err != nil {
+		return fmt.Errorf("error setting direct_internet_access for sagemaker notebook instance (%s): %s", d.Id(), err)
+	}
+
 	tagsOutput, err := conn.ListTags(&sagemaker.ListTagsInput{
 		ResourceArn: notebookInstance.NotebookInstanceArn,
 	})

aws/resource_aws_sagemaker_notebook_instance_test.go

@@ -324,6 +324,57 @@ func testAccCheckAWSSagemakerNotebookInstanceName(notebook *sagemaker.DescribeNo
 	}
 }
 
+func TestAccAWSSagemakerNotebookInstance_direct_internet_access(t *testing.T) {
+	var notebook sagemaker.DescribeNotebookInstanceOutput
+	notebookName := resource.PrefixedUniqueId(sagemakerTestAccSagemakerNotebookInstanceResourceNamePrefix)
+	var resourceName = "aws_sagemaker_notebook_instance.foo"
+	resource.ParallelTest(t, resource.TestCase{
+		PreCheck:     func() { testAccPreCheck(t) },
+		Providers:    testAccProviders,
+		CheckDestroy: testAccCheckAWSSagemakerNotebookInstanceDestroy,
+		Steps: []resource.TestStep{
+			{
+				Config: testAccAWSSagemakerNotebookInstanceConfigDirectInternetAccess(notebookName, "Disabled"),
+				Check: resource.ComposeTestCheckFunc(
+					testAccCheckAWSSagemakerNotebookInstanceExists(resourceName, &notebook),
+					testAccCheckAWSSagemakerNotebookInstanceName(&notebook, notebookName),
+					testAccCheckAWSSagemakerNotebookDirectInternetAccess(&notebook, "Disabled"),
+
+					resource.TestCheckResourceAttr(
+						"aws_sagemaker_notebook_instance.foo", "name", notebookName),
+				),
+			},
+			{
+				Config: testAccAWSSagemakerNotebookInstanceConfigDirectInternetAccess(notebookName, "Enabled"),
+				Check: resource.ComposeTestCheckFunc(
+					testAccCheckAWSSagemakerNotebookInstanceExists(resourceName, &notebook),
+					testAccCheckAWSSagemakerNotebookInstanceName(&notebook, notebookName),
+					testAccCheckAWSSagemakerNotebookDirectInternetAccess(&notebook, "Enabled"),
+
+					resource.TestCheckResourceAttr(
+						"aws_sagemaker_notebook_instance.foo", "name", notebookName),
+				),
+			},
+			{
+				ResourceName:      resourceName,
+				ImportState:       true,
+				ImportStateVerify: true,
+			},
+		},
+	})
+}
+
+func testAccCheckAWSSagemakerNotebookDirectInternetAccess(notebook *sagemaker.DescribeNotebookInstanceOutput, expected string) resource.TestCheckFunc {
+	return func(s *terraform.State) error {
+		directInternetAccess := notebook.DirectInternetAccess
+		if *directInternetAccess != expected {
+			return fmt.Errorf("direct_internet_access setting is incorrect: %s", *notebook.DirectInternetAccess)
+		}
+
+		return nil
+	}
+}
+
 func testAccCheckAWSSagemakerNotebookInstanceTags(notebook *sagemaker.DescribeNotebookInstanceOutput, key string, value string) resource.TestCheckFunc {
 	return func(s *terraform.State) error {
 		conn := testAccProvider.Meta().(*AWSClient).sagemakerconn
@@ -501,3 +552,48 @@ data "aws_iam_policy_document" "assume_role" {
 }
 `, notebookName, notebookName)
 }
+
+func testAccAWSSagemakerNotebookInstanceConfigDirectInternetAccess(notebookName string, directInternetAccess string) string {
+	return fmt.Sprintf(`
+resource "aws_sagemaker_notebook_instance" "foo" {
+	name = %[1]q
+	role_arn = "${aws_iam_role.foo.arn}"
+	instance_type = "ml.t2.medium"
+	subnet_id = "${aws_subnet.sagemaker.id}"
+	direct_internet_access = %[2]q
+}
+
+resource "aws_iam_role" "foo" {
+	name = %[1]q
+	path = "/"
+	assume_role_policy = "${data.aws_iam_policy_document.assume_role.json}"
+}
+
+data "aws_iam_policy_document" "assume_role" {
+	statement {
+		actions = [ "sts:AssumeRole" ]
+		principals {
+			type = "Service"
+			identifiers = [ "sagemaker.amazonaws.com" ]
+		}
+	}
+}
+
+resource "aws_vpc" "test" {
+  cidr_block = "10.0.10.0/24"
+
+  tags = {
+    Name = "tf-acc-test-sagemaker-notebook-instance-direct-internet-access"
+  }
+}
+
+resource "aws_subnet" "sagemaker" {
+  vpc_id     = "${aws_vpc.test.id}"
+  cidr_block = "10.0.0.0/27"
+
+  tags = {
+    Name = "tf-acc-test-sagemaker-notebook-instance-direct-internet-access"
+  }
+}
+`, notebookName, directInternetAccess)
+}

website/docs/r/sagemaker_notebook_instance.html.markdown

@@ -37,6 +37,7 @@ The following arguments are supported:
 * `security_groups` - (Optional) The associated security groups.
 * `kms_key_id` - (Optional) The AWS Key Management Service (AWS KMS) key that Amazon SageMaker uses to encrypt the model artifacts at rest using Amazon S3 server-side encryption.
 *  `lifecycle_config_name` - (Optional) The name of a lifecycle configuration to associate with the notebook instance.
+* `direct_internet_access` - (Optional) Set to `Disabled` to disable internet access to notebook. Requires `subnet_id` to be set. Supported values: `Enabled`(Default) or `Disabled`. If set to `Disabled`, the notebook instance will be able to access resources only in your VPC, and will not be able to connect to Amazon SageMaker training and endpoint services unless your configure a NAT Gateway in your VPC. 
 * `tags` - (Optional) A mapping of tags to assign to the resource.
 
 ## Attributes Reference