Merge pull request #40771 from acwwat/f-aws_emr_studio-add_enryption_…

…key_arn_arg

feat: Add encryption_key_arn arg for aws_emr_studio

.changelog/40771.txt

@@ -0,0 +1,7 @@
+```release-note:enhancement
+resource/aws_emr_studio: Add `encryption_key_arn` argument
+```
+
+```release-note:bug
+resource/aws_emr_studio: Fix issue with IAM/KMS policy eventual consistency handling not working
+```

internal/service/emr/studio.go

@@ -60,6 +60,11 @@ func resourceStudio() *schema.Resource {
 				Optional:     true,
 				ValidateFunc: validation.StringLenBetween(0, 256),
 			},
+			"encryption_key_arn": {
+				Type:         schema.TypeString,
+				Optional:     true,
+				ValidateFunc: verify.ValidARN,
+			},
 			"engine_security_group_id": {
 				Type:     schema.TypeString,
 				Required: true,
@@ -140,6 +145,10 @@ func resourceStudioCreate(ctx context.Context, d *schema.ResourceData, meta inte
 		input.Description = aws.String(v.(string))
 	}
 
+	if v, ok := d.GetOk("encryption_key_arn"); ok {
+		input.EncryptionKeyArn = aws.String(v.(string))
+	}
+
 	if v, ok := d.GetOk("idp_auth_url"); ok {
 		input.IdpAuthUrl = aws.String(v.(string))
 	}
@@ -158,7 +167,8 @@ func resourceStudioCreate(ctx context.Context, d *schema.ResourceData, meta inte
 		},
 		func(err error) (bool, error) {
 			if errs.IsAErrorMessageContains[*awstypes.InvalidRequestException](err, "entity does not have permissions to assume role") ||
-				errs.IsAErrorMessageContains[*awstypes.InvalidRequestException](err, "Service role does not have permission to access") {
+				errs.IsAErrorMessageContains[*awstypes.InvalidRequestException](err, "Service role does not have permission to access") ||
+				errs.IsAErrorMessageContains[*awstypes.InvalidRequestException](err, "'ServiceRole' does not have permission to access") {
 				return true, err
 			}
 
@@ -195,6 +205,7 @@ func resourceStudioRead(ctx context.Context, d *schema.ResourceData, meta interf
 	d.Set("auth_mode", studio.AuthMode)
 	d.Set("default_s3_location", studio.DefaultS3Location)
 	d.Set(names.AttrDescription, studio.Description)
+	d.Set("encryption_key_arn", studio.EncryptionKeyArn)
 	d.Set("engine_security_group_id", studio.EngineSecurityGroupId)
 	d.Set("idp_auth_url", studio.IdpAuthUrl)
 	d.Set("idp_relay_state_parameter_name", studio.IdpRelayStateParameterName)
@@ -228,6 +239,10 @@ func resourceStudioUpdate(ctx context.Context, d *schema.ResourceData, meta inte
 			input.Description = aws.String(d.Get(names.AttrDescription).(string))
 		}
 
+		if d.HasChange("encryption_key_arn") {
+			input.EncryptionKeyArn = aws.String(d.Get("encryption_key_arn").(string))
+		}
+
 		if d.HasChange(names.AttrName) {
 			input.Name = aws.String(d.Get(names.AttrName).(string))
 		}

internal/service/emr/studio_test.go

@@ -47,6 +47,7 @@ func TestAccEMRStudio_sso(t *testing.T) {
 					resource.TestCheckResourceAttrPair(resourceName, "user_role", "aws_iam_role.test", names.AttrARN),
 					resource.TestCheckResourceAttrPair(resourceName, names.AttrServiceRole, "aws_iam_role.test", names.AttrARN),
 					resource.TestCheckResourceAttr(resourceName, "subnet_ids.#", "1"),
+					resource.TestCheckResourceAttr(resourceName, "encryption_key_arn", ""),
 					resource.TestCheckResourceAttr(resourceName, acctest.CtTagsPercent, "0"),
 				),
 			},
@@ -69,6 +70,7 @@ func TestAccEMRStudio_sso(t *testing.T) {
 					resource.TestCheckResourceAttrPair(resourceName, "user_role", "aws_iam_role.test", names.AttrARN),
 					resource.TestCheckResourceAttrPair(resourceName, names.AttrServiceRole, "aws_iam_role.test", names.AttrARN),
 					resource.TestCheckResourceAttr(resourceName, "subnet_ids.#", "1"),
+					resource.TestCheckResourceAttr(resourceName, "encryption_key_arn", ""),
 					resource.TestCheckResourceAttr(resourceName, acctest.CtTagsPercent, "0"),
 				),
 			},
@@ -100,6 +102,7 @@ func TestAccEMRStudio_iam(t *testing.T) {
 					resource.TestCheckResourceAttrPair(resourceName, "engine_security_group_id", "aws_security_group.test", names.AttrID),
 					resource.TestCheckResourceAttrPair(resourceName, names.AttrServiceRole, "aws_iam_role.test", names.AttrARN),
 					resource.TestCheckResourceAttr(resourceName, "subnet_ids.#", "1"),
+					resource.TestCheckResourceAttr(resourceName, "encryption_key_arn", ""),
 					resource.TestCheckResourceAttr(resourceName, acctest.CtTagsPercent, "0"),
 				),
 			},
@@ -137,6 +140,43 @@ func TestAccEMRStudio_disappears(t *testing.T) {
 	})
 }
 
+func TestAccEMRStudio_workspaceStorageEncryption(t *testing.T) {
+	ctx := acctest.Context(t)
+	var studio awstypes.Studio
+	resourceName := "aws_emr_studio.test"
+	rName := sdkacctest.RandomWithPrefix(acctest.ResourcePrefix)
+
+	resource.ParallelTest(t, resource.TestCase{
+		PreCheck:                 func() { acctest.PreCheck(ctx, t) },
+		ErrorCheck:               acctest.ErrorCheck(t, names.EMRServiceID),
+		ProtoV5ProviderFactories: acctest.ProtoV5ProviderFactories,
+		CheckDestroy:             testAccCheckStudioDestroy(ctx),
+		Steps: []resource.TestStep{
+			{
+				Config: testAccStudioConfig_workspaceStorageEncryption(rName),
+				Check: resource.ComposeTestCheckFunc(
+					testAccCheckStudioExists(ctx, resourceName, &studio),
+					resource.TestCheckResourceAttr(resourceName, names.AttrName, rName),
+					resource.TestCheckResourceAttr(resourceName, "auth_mode", "IAM"),
+					resource.TestCheckResourceAttrSet(resourceName, names.AttrURL),
+					resource.TestCheckResourceAttrPair(resourceName, names.AttrVPCID, "aws_vpc.test", names.AttrID),
+					resource.TestCheckResourceAttrPair(resourceName, "workspace_security_group_id", "aws_security_group.test", names.AttrID),
+					resource.TestCheckResourceAttrPair(resourceName, "engine_security_group_id", "aws_security_group.test", names.AttrID),
+					resource.TestCheckResourceAttrPair(resourceName, names.AttrServiceRole, "aws_iam_role.test", names.AttrARN),
+					resource.TestCheckResourceAttr(resourceName, "subnet_ids.#", "1"),
+					resource.TestCheckResourceAttrPair(resourceName, "encryption_key_arn", "aws_kms_key.test", names.AttrARN),
+					resource.TestCheckResourceAttr(resourceName, acctest.CtTagsPercent, "0"),
+				),
+			},
+			{
+				ResourceName:      resourceName,
+				ImportState:       true,
+				ImportStateVerify: true,
+			},
+		},
+	})
+}
+
 func TestAccEMRStudio_tags(t *testing.T) {
 	ctx := acctest.Context(t)
 	var studio awstypes.Studio
@@ -245,6 +285,15 @@ resource "aws_iam_role" "test" {
   assume_role_policy = data.aws_iam_policy_document.assume_role.json
 }
 
+data "aws_iam_policy" "emr" {
+  name = "AmazonEMRServicePolicy_v2"
+}
+
+resource "aws_iam_role_policy_attachment" "test-attach" {
+  role       = aws_iam_role.test.name
+  policy_arn = data.aws_iam_policy.emr.arn
+}
+
 data "aws_iam_policy_document" "assume_role" {
   statement {
     actions = ["sts:AssumeRole"]
@@ -321,16 +370,102 @@ resource "aws_emr_studio" "test" {
 `, rName))
 }
 
+func testAccStudioConfig_workspaceStorageEncryption(rName string) string {
+	return acctest.ConfigCompose(testAccStudioConfig_base(rName), fmt.Sprintf(`
+data "aws_caller_identity" "current" {}
+
+data "aws_region" "current" {}
+
+resource "aws_kms_key" "test" {
+  description = %[1]q
+
+  policy = <<POLICY
+{
+  "Version": "2012-10-17",
+  "Id": "kms-tf-1",
+  "Statement": [
+    {
+      "Sid": "Enable IAM User Permissions",
+      "Effect": "Allow",
+      "Principal": {
+        "AWS": "*"
+      },
+      "Action": "kms:*",
+      "Resource": "*"
+    },
+    {
+      "Sid": "AllowEMRStudioServiceRoleAccess",
+      "Effect": "Allow",
+      "Principal": {
+        "AWS": "${aws_iam_role.test.arn}"
+      },
+      "Action": [
+        "kms:Decrypt",
+        "kms:GenerateDataKey",
+        "kms:ReEncryptFrom",
+        "kms:ReEncryptTo",
+        "kms:DescribeKey"
+      ],
+      "Resource": "*",
+      "Condition": {
+        "StringEquals": {
+          "kms:CallerAccount": "${data.aws_caller_identity.current.account_id}",
+          "kms:EncryptionContext:aws:s3:arn": "${aws_s3_bucket.test.arn}",
+          "kms:ViaService": "s3.${data.aws_region.current.name}.amazonaws.com"
+        }
+      }
+    }
+  ]
+}
+POLICY
+}
+
+resource "aws_iam_role_policy" "test_kms" {
+  name   = "%[1]s-kms"
+  role   = aws_iam_role.test.id
+  policy = <<EOF
+{
+  "Version": "2012-10-17",
+  "Statement": [
+    {
+      "Effect": "Allow",
+      "Action": [
+        "kms:Decrypt",
+        "kms:GenerateDataKey",
+        "kms:ReEncryptFrom",
+        "kms:ReEncryptTo",
+        "kms:DescribeKey"
+      ],
+      "Resource": "${aws_kms_key.test.arn}"
+    }
+  ]
+}
+EOF
+}
+
+resource "aws_emr_studio" "test" {
+  auth_mode                   = "IAM"
+  default_s3_location         = "s3://${aws_s3_bucket.test.bucket}/test"
+  encryption_key_arn          = aws_kms_key.test.arn
+  engine_security_group_id    = aws_security_group.test.id
+  name                        = %[1]q
+  service_role                = aws_iam_role.test.arn
+  subnet_ids                  = aws_subnet.test[*].id
+  vpc_id                      = aws_vpc.test.id
+  workspace_security_group_id = aws_security_group.test.id
+}
+`, rName))
+}
+
 func testAccStudioConfig_tags1(rName, tagKey1, tagValue1 string) string {
 	return acctest.ConfigCompose(testAccStudioConfig_base(rName), fmt.Sprintf(`
 resource "aws_emr_studio" "test" {
-  auth_mode                   = "SSO"
+  auth_mode                   = "IAM"
   default_s3_location         = "s3://${aws_s3_bucket.test.bucket}/test"
   engine_security_group_id    = aws_security_group.test.id
   name                        = %[1]q
   service_role                = aws_iam_role.test.arn
   subnet_ids                  = aws_subnet.test[*].id
-  user_role                   = aws_iam_role.test.arn
   vpc_id                      = aws_vpc.test.id
   workspace_security_group_id = aws_security_group.test.id
 
@@ -344,13 +479,12 @@ resource "aws_emr_studio" "test" {
 func testAccStudioConfig_tags2(rName, tagKey1, tagValue1, tagKey2, tagValue2 string) string {
 	return acctest.ConfigCompose(testAccStudioConfig_base(rName), fmt.Sprintf(`
 resource "aws_emr_studio" "test" {
-  auth_mode                   = "SSO"
+  auth_mode                   = "IAM"
   default_s3_location         = "s3://${aws_s3_bucket.test.bucket}/test"
   engine_security_group_id    = aws_security_group.test.id
   name                        = %[1]q
   service_role                = aws_iam_role.test.arn
   subnet_ids                  = aws_subnet.test[*].id
-  user_role                   = aws_iam_role.test.arn
   vpc_id                      = aws_vpc.test.id
   workspace_security_group_id = aws_security_group.test.id
 

website/docs/r/emr_studio.html.markdown

@@ -42,6 +42,7 @@ The following arguments are required:
 The following arguments are optional:
 
 * `description` - (Optional) A detailed description of the Amazon EMR Studio.
+* `encryption_key_arn` - (Optional) The AWS KMS key identifier (ARN) used to encrypt Amazon EMR Studio workspace and notebook files when backed up to Amazon S3.
 * `idp_auth_url` - (Optional) The authentication endpoint of your identity provider (IdP). Specify this value when you use IAM authentication and want to let federated users log in to a Studio with the Studio URL and credentials from your IdP. Amazon EMR Studio redirects users to this endpoint to enter credentials.
 * `idp_relay_state_parameter_name` - (Optional) The name that your identity provider (IdP) uses for its RelayState parameter. For example, RelayState or TargetSource. Specify this value when you use IAM authentication and want to let federated users log in to a Studio using the Studio URL. The RelayState parameter differs by IdP.
 * `tags` - (Optional) list of tags to apply to the EMR Cluster. If configured with a provider [`default_tags` configuration block](https://registry.terraform.io/providers/hashicorp/aws/latest/docs#default_tags-configuration-block) present, tags with matching keys will overwrite those defined at the provider-level.