Resolve cloudfrontkeyvaluestore_keys_exclusive being blocked by service quotas with large number of key value pairs

.changelog/42795.txt

@@ -0,0 +1,2 @@
+```release-note:bug
+resource/aws_cloudfrontkeyvaluestore_keys_exclusive: Ensure that UpdateKeys API calls are batched to stay under the Key Value Store Service Quota. The `max_batch_size` schema can be used to override the default when a user has key value store items that are large in size and need to use smaller batches.

internal/service/cloudfrontkeyvaluestore/keys_exclusive.go

@@ -5,16 +5,21 @@ package cloudfrontkeyvaluestore
 
 import (
 	"context"
+	"slices"
 
 	"github.com/aws/aws-sdk-go-v2/aws"
 	"github.com/aws/aws-sdk-go-v2/service/cloudfrontkeyvaluestore"
 	awstypes "github.com/aws/aws-sdk-go-v2/service/cloudfrontkeyvaluestore/types"
+	"github.com/hashicorp/terraform-plugin-framework-validators/int64validator"
 	"github.com/hashicorp/terraform-plugin-framework/diag"
 	"github.com/hashicorp/terraform-plugin-framework/path"
 	"github.com/hashicorp/terraform-plugin-framework/resource"
 	"github.com/hashicorp/terraform-plugin-framework/resource/schema"
+	"github.com/hashicorp/terraform-plugin-framework/resource/schema/int64default"
+	"github.com/hashicorp/terraform-plugin-framework/resource/schema/int64planmodifier"
 	"github.com/hashicorp/terraform-plugin-framework/resource/schema/planmodifier"
 	"github.com/hashicorp/terraform-plugin-framework/resource/schema/stringplanmodifier"
+	"github.com/hashicorp/terraform-plugin-framework/schema/validator"
 	"github.com/hashicorp/terraform-plugin-framework/types"
 	"github.com/hashicorp/terraform-plugin-sdk/v2/helper/retry"
 	"github.com/hashicorp/terraform-provider-aws/internal/conns"
@@ -32,6 +37,7 @@ import (
 const (
 	ResNameKeysExclusive = "Keys Exclusive"
 	ResNameKeyValueStore = "Key Value Store"
+	maxBatchSizeDefault  = 50
 )
 
 // @FrameworkResource("aws_cloudfrontkeyvaluestore_keys_exclusive", name="Keys  Exclusive")
@@ -55,6 +61,18 @@ func (r *resourceKeysExclusive) Schema(ctx context.Context, request resource.Sch
 					stringplanmodifier.RequiresReplace(),
 				},
 			},
+			"max_batch_size": schema.Int64Attribute{
+				Optional: true,
+				Computed: true,
+				Default:  int64default.StaticInt64(maxBatchSizeDefault),
+				Validators: []validator.Int64{
+					int64validator.Between(1, 50),
+				},
+				PlanModifiers: []planmodifier.Int64{
+					int64planmodifier.UseStateForUnknown(),
+				},
+				MarkdownDescription: "Maximum resource key values pairs that you wills update in a single API request. AWS has a default quota of 50 keys or a 3 MB payload, whichever is reached first",
+			},
 			"total_size_in_bytes": schema.Int64Attribute{
 				Computed:            true,
 				MarkdownDescription: "Total size of the Key Value Store in bytes.",
@@ -108,38 +126,54 @@ func (r *resourceKeysExclusive) syncKeyValuePairs(ctx context.Context, plan *res
 
 	put, del, _ := intflex.DiffSlices(have, want, resourceKeyValuePairEqual)
 
-	putRequired := len(put) > 0
-	deleteRequired := len(del) > 0
+	// We need to perform a batched operation in the event of many Key Value Pairs
+	// to stay within AWS service limits
+	//
+	// https://docs.aws.amazon.com/AmazonCloudFront/latest/DeveloperGuide/cloudfront-limits.html#limits-keyvaluestores
+	batchSize := int(plan.MaximumBatchSize.ValueInt64())
+	etag := kvs.ETag
+	totalSizeInBytes := kvs.TotalSizeInBytes
 
-	if putRequired || deleteRequired {
+	for chunk := range slices.Chunk(expandPutKeyRequestListItem(put), batchSize) {
 		input := cloudfrontkeyvaluestore.UpdateKeysInput{
 			KvsARN:  aws.String(kvsARN),
-			IfMatch: kvs.ETag,
+			IfMatch: etag,
+			Puts:    chunk,
 		}
 
-		if putRequired {
-			input.Puts = expandPutKeyRequestListItem(put)
+		out, err := conn.UpdateKeys(ctx, &input)
+		if err != nil {
+			diags.AddError(
+				create.ProblemStandardMessage(names.CloudFrontKeyValueStore, create.ErrActionSynchronizing, ResNameKeysExclusive, kvsARN, err),
+				err.Error(),
+			)
+			return diags
 		}
+		etag = out.ETag
+		totalSizeInBytes = out.TotalSizeInBytes
+	}
 
-		if deleteRequired {
-			input.Deletes = expandDeleteKeyRequestListItem(del)
+	for chunk := range slices.Chunk(expandDeleteKeyRequestListItem(del), batchSize) {
+		input := cloudfrontkeyvaluestore.UpdateKeysInput{
+			KvsARN:  aws.String(kvsARN),
+			IfMatch: etag,
+			Deletes: chunk,
 		}
 
 		out, err := conn.UpdateKeys(ctx, &input)
-
 		if err != nil {
 			diags.AddError(
 				create.ProblemStandardMessage(names.CloudFrontKeyValueStore, create.ErrActionSynchronizing, ResNameKeysExclusive, kvsARN, err),
 				err.Error(),
 			)
 			return diags
 		}
-
-		plan.TotalSizeInBytes = flex.Int64ToFramework(ctx, out.TotalSizeInBytes)
-	} else {
-		plan.TotalSizeInBytes = flex.Int64ToFramework(ctx, kvs.TotalSizeInBytes)
+		etag = out.ETag
+		totalSizeInBytes = out.TotalSizeInBytes
 	}
 
+	plan.TotalSizeInBytes = flex.Int64ToFramework(ctx, totalSizeInBytes)
+
 	return diags
 }
 
@@ -185,6 +219,10 @@ func (r *resourceKeysExclusive) Read(ctx context.Context, request resource.ReadR
 	data.KvsARN = fwtypes.ARNValue(aws.ToString(kvs.KvsARN))
 	data.TotalSizeInBytes = types.Int64Value(aws.ToInt64(kvs.TotalSizeInBytes))
 
+	if data.MaximumBatchSize.IsNull() || data.MaximumBatchSize.ValueInt64() == 0 {
+		data.MaximumBatchSize = types.Int64Value(maxBatchSizeDefault)
+	}
+
 	response.Diagnostics.Append(flex.Flatten(ctx, keyPairs, &data.ResourceKeyValuePair)...)
 	if response.Diagnostics.HasError() {
 		return
@@ -300,5 +338,6 @@ type resourceKeyValuePairModel struct {
 type resourceKeysExclusiveModel struct {
 	ResourceKeyValuePair fwtypes.SetNestedObjectValueOf[resourceKeyValuePairModel] `tfsdk:"resource_key_value_pair"`
 	KvsARN               fwtypes.ARN                                               `tfsdk:"key_value_store_arn"`
+	MaximumBatchSize     types.Int64                                               `tfsdk:"max_batch_size"`
 	TotalSizeInBytes     types.Int64                                               `tfsdk:"total_size_in_bytes"`
 }

internal/service/cloudfrontkeyvaluestore/keys_exclusive_test.go

@@ -35,7 +35,9 @@ func TestAccCloudFrontKeyValueStoreKeysExclusive_basic(t *testing.T) {
 	rName := sdkacctest.RandomWithPrefix(acctest.ResourcePrefix)
 	var keys []string
 	var values []string
-	for i := 1; i < 6; i++ {
+
+	// Test with a large number of key value pairs to ensure batching is working correctly
+	for i := 1; i < 170; i++ {
 		keys = append(keys, sdkacctest.RandomWithPrefix(acctest.ResourcePrefix))
 		values = append(values, sdkacctest.RandomWithPrefix(acctest.ResourcePrefix))
 	}
@@ -287,6 +289,50 @@ func TestAccCloudFrontKeyValueStoreKeysExclusive_empty(t *testing.T) {
 	})
 }
 
+func TestAccCloudFrontKeyValueStoreKeysExclusive_maxBatchSize(t *testing.T) {
+	ctx := acctest.Context(t)
+	rName := sdkacctest.RandomWithPrefix(acctest.ResourcePrefix)
+	maxBatchSize := sdkacctest.RandIntRange(35, 49)
+	var keys []string
+	var values []string
+	// Test with a large number of key value pairs to ensure batching is working correctly
+	for i := 1; i < 170; i++ {
+		keys = append(keys, sdkacctest.RandomWithPrefix(acctest.ResourcePrefix))
+		values = append(values, sdkacctest.RandomWithPrefix(acctest.ResourcePrefix))
+	}
+	resourceName := "aws_cloudfrontkeyvaluestore_keys_exclusive.test"
+	kvsResourceName := "aws_cloudfront_key_value_store.test"
+
+	resource.ParallelTest(t, resource.TestCase{
+		PreCheck: func() {
+			acctest.PreCheck(ctx, t)
+			acctest.PreCheckPartitionHasService(t, names.CloudFront)
+		},
+		ErrorCheck:               acctest.ErrorCheck(t, names.CloudFront),
+		ProtoV5ProviderFactories: acctest.ProtoV5ProviderFactories,
+		CheckDestroy:             testAccCheckKeysExclusiveDestroy(ctx),
+		Steps: []resource.TestStep{
+			{
+				Config: testAccKeysExclusiveConfig_maxBatchSize(keys, values, rName, maxBatchSize),
+				Check: resource.ComposeTestCheckFunc(
+					testAccCheckKeysExclusiveExists(ctx, resourceName),
+					resource.TestCheckResourceAttrPair(resourceName, "key_value_store_arn", kvsResourceName, names.AttrARN),
+					resource.TestCheckResourceAttrSet(resourceName, "total_size_in_bytes"),
+					testCheckMultipleKeyValuePairs(keys, values, resourceName),
+				),
+			},
+			{
+				ResourceName:                         resourceName,
+				ImportState:                          true,
+				ImportStateIdFunc:                    acctest.AttrImportStateIdFunc(resourceName, "key_value_store_arn"),
+				ImportStateVerify:                    true,
+				ImportStateVerifyIdentifierAttribute: "key_value_store_arn",
+				ImportStateVerifyIgnore:              []string{"max_batch_size"},
+			},
+		},
+	})
+}
+
 func testAccCheckKeysExclusiveDestroy(ctx context.Context) resource.TestCheckFunc {
 	return func(s *terraform.State) error {
 		conn := acctest.Provider.Meta().(*conns.AWSClient).CloudFrontKeyValueStoreClient(ctx)
@@ -431,3 +477,34 @@ resource "aws_cloudfrontkeyvaluestore_keys_exclusive" "test" {
 }
 `, rName)
 }
+
+func testAccKeysExclusiveConfig_maxBatchSize(keys, values []string, rName string, maxBatchSize int) string {
+	keysJson, _ := json.Marshal(keys)
+	keysString := string(keysJson)
+	valuesJson, _ := json.Marshal(values)
+	valuesString := string(valuesJson)
+	return fmt.Sprintf(`
+locals {
+  key_list      = %[1]s
+  value_list    = %[2]s
+  key_value_set = { for i, v in local.key_list : local.key_list[i] => local.value_list[i] }
+}
+
+resource "aws_cloudfront_key_value_store" "test" {
+  name = %[3]q
+}
+
+resource "aws_cloudfrontkeyvaluestore_keys_exclusive" "test" {
+  key_value_store_arn = aws_cloudfront_key_value_store.test.arn
+  max_batch_size      = %[4]d
+  dynamic "resource_key_value_pair" {
+    for_each = local.key_value_set
+    content {
+      key   = resource_key_value_pair.key
+      value = resource_key_value_pair.value
+
+    }
+  }
+}
+`, keysString, valuesString, rName, maxBatchSize)
+}

website/docs/r/cloudfrontkeyvaluestore_keys_exclusive.html.markdown

@@ -53,6 +53,7 @@ The following arguments are required:
 
 The following arguments are optional:
 
+* `max_batch_size` - (Optional) Maximum resource key values pairs that will update in a single API request. AWS has a default quota of 50 keys or a 3 MB payload, whichever is reached first. Defaults to `50`.
 * `resource_key_value_pair` - (Optional) A list of all resource key value pairs associated with the KeyValueStore.
 See [`resource_key_value_pair`](#resource_key_value_pair) below.