Allows secondary_private_ip_address_count to be updated in place

.changelog/47477.txt

@@ -0,0 +1,3 @@
+```release-note:bug
+resource/aws_nat_gateway: Allow updating `secondary_private_ip_address_count` in-place for private NAT gateways
+```

internal/service/ec2/status.go

@@ -1725,6 +1725,39 @@ func statusNATGatewayAddressByNATGatewayIDAndPrivateIP(conn *ec2.Client, natGate
 	}
 }
 
+func statusNATGatewaySecondaryPrivateIPAddressCount(conn *ec2.Client, natGatewayID string, expectedCount int) retry.StateRefreshFunc {
+	return func(ctx context.Context) (any, string, error) {
+		output, err := findNATGatewayByID(ctx, conn, natGatewayID)
+
+		if retry.NotFound(err) {
+			return nil, "", nil
+		}
+
+		if err != nil {
+			return nil, "", err
+		}
+
+		secondaryCount := 0
+		for _, natGatewayAddress := range output.NatGatewayAddresses {
+			if !aws.ToBool(natGatewayAddress.IsPrimary) && aws.ToString(natGatewayAddress.PrivateIp) != "" {
+				if natGatewayAddress.Status == awstypes.NatGatewayAddressStatusFailed {
+					return output, string(natGatewayAddress.Status), nil
+				}
+
+				if natGatewayAddress.Status == awstypes.NatGatewayAddressStatusSucceeded {
+					secondaryCount++
+				}
+			}
+		}
+
+		if secondaryCount == expectedCount {
+			return output, "ready", nil
+		}
+
+		return output, "pending", nil
+	}
+}
+
 func statusNATGatewayAttachedAppliances(conn *ec2.Client, id string) retry.StateRefreshFunc {
 	return func(ctx context.Context) (any, string, error) {
 		output, err := findNATGatewayByID(ctx, conn, id)

internal/service/ec2/vpc_nat_gateway.go

@@ -432,7 +432,12 @@ func resourceNATGatewayUpdate(ctx context.Context, d *schema.ResourceData, meta
 	case awstypes.AvailabilityModeZonal:
 		switch awstypes.ConnectivityType(d.Get("connectivity_type").(string)) {
 		case awstypes.ConnectivityTypePrivate:
-			if d.HasChanges("secondary_private_ip_addresses") {
+			countRaw := d.GetRawConfig().GetAttr("secondary_private_ip_address_count")
+			countConfigured := countRaw.IsKnown() && !countRaw.IsNull()
+			addressesRaw := d.GetRawConfig().GetAttr("secondary_private_ip_addresses")
+			addressesConfigured := addressesRaw.IsKnown() && !addressesRaw.IsNull()
+
+			if addressesConfigured && d.HasChange("secondary_private_ip_addresses") {
 				o, n := d.GetChange("secondary_private_ip_addresses")
 				os, ns := o.(*schema.Set), n.(*schema.Set)
 
@@ -474,6 +479,57 @@ func resourceNATGatewayUpdate(ctx context.Context, d *schema.ResourceData, meta
 					}
 				}
 			}
+
+			if countConfigured && d.HasChange("secondary_private_ip_address_count") {
+				o, n := d.GetChange("secondary_private_ip_address_count")
+				oldCount, newCount := o.(int), n.(int)
+
+				delta := newCount - oldCount
+
+				if delta > 0 {
+					input := &ec2.AssignPrivateNatGatewayAddressInput{
+						NatGatewayId:          aws.String(d.Id()),
+						PrivateIpAddressCount: aws.Int32(int32(delta)),
+					}
+
+					_, err := conn.AssignPrivateNatGatewayAddress(ctx, input)
+
+					if err != nil {
+						return sdkdiag.AppendErrorf(diags, "assigning EC2 NAT Gateway (%s) private IP address count: %s", d.Id(), err)
+					}
+
+					if _, err := waitNATGatewaySecondaryPrivateIPAddressCount(ctx, conn, d.Id(), newCount, d.Timeout(schema.TimeoutUpdate)); err != nil {
+						return sdkdiag.AppendErrorf(diags, "waiting for EC2 NAT Gateway (%s) secondary private IP address count (%d): %s", d.Id(), newCount, err)
+					}
+				}
+
+				if delta < 0 {
+					removeCount := -delta
+
+					sIPs, _ := d.GetChange("secondary_private_ip_addresses")
+					secondaryPrivateIPs := sIPs.(*schema.Set).List()
+
+					privateIPsToUnassign := secondaryPrivateIPs[:removeCount]
+
+					input := &ec2.UnassignPrivateNatGatewayAddressInput{
+						NatGatewayId:            aws.String(d.Id()),
+						PrivateIpAddresses:      flex.ExpandStringValueList(privateIPsToUnassign),
+						MaxDrainDurationSeconds: aws.Int32(50),
+					}
+
+					_, err := conn.UnassignPrivateNatGatewayAddress(ctx, input)
+
+					if err != nil {
+						return sdkdiag.AppendErrorf(diags, "unassigning EC2 NAT Gateway (%s) private IP addresses: %s", d.Id(), err)
+					}
+
+					for _, privateIP := range flex.ExpandStringValueList(privateIPsToUnassign) {
+						if _, err := waitNATGatewayAddressUnassigned(ctx, conn, d.Id(), privateIP, d.Timeout(schema.TimeoutUpdate)); err != nil {
+							return sdkdiag.AppendErrorf(diags, "waiting for EC2 NAT Gateway (%s) private IP address (%s) unassign: %s", d.Id(), privateIP, err)
+						}
+					}
+				}
+			}
 		case awstypes.ConnectivityTypePublic:
 			if !d.GetRawConfig().GetAttr("secondary_allocation_ids").IsNull() && d.HasChanges("secondary_allocation_ids") {
 				o, n := d.GetChange("secondary_allocation_ids")
@@ -654,15 +710,18 @@ func resourceNATGatewayCustomizeDiff(ctx context.Context, diff *schema.ResourceD
 			return fmt.Errorf(`secondary_allocation_ids is not supported with connectivity_type = "%s"`, connectivityType)
 		}
 
-		if diff.Id() != "" && diff.HasChange("secondary_private_ip_address_count") {
-			if v := diff.GetRawConfig().GetAttr("secondary_private_ip_address_count"); v.IsKnown() && !v.IsNull() {
-				if err := diff.ForceNew("secondary_private_ip_address_count"); err != nil {
-					return fmt.Errorf("setting secondary_private_ip_address_count to ForceNew: %w", err)
-				}
+		countRaw := diff.GetRawConfig().GetAttr("secondary_private_ip_address_count")
+		countConfigured := countRaw.IsKnown() && !countRaw.IsNull()
+		addressesRaw := diff.GetRawConfig().GetAttr("secondary_private_ip_addresses")
+		addressesConfigured := addressesRaw.IsKnown() && !addressesRaw.IsNull()
+
+		if diff.Id() != "" && countConfigured && diff.HasChange("secondary_private_ip_address_count") {
+			if err := diff.SetNewComputed("secondary_private_ip_addresses"); err != nil {
+				return fmt.Errorf("setting secondary_private_ip_addresses to Computed: %w", err)
 			}
 		}
 
-		if diff.Id() != "" && diff.HasChange("secondary_private_ip_addresses") {
+		if diff.Id() != "" && addressesConfigured && diff.HasChange("secondary_private_ip_addresses") {
 			if err := diff.SetNewComputed("secondary_private_ip_address_count"); err != nil {
 				return fmt.Errorf("setting secondary_private_ip_address_count to Computed: %w", err)
 			}

internal/service/ec2/vpc_nat_gateway_test.go

@@ -430,6 +430,62 @@ func TestAccVPCNATGateway_secondaryPrivateIPAddressCountToSpecific(t *testing.T)
 	})
 }
 
+func TestAccVPCNATGateway_secondaryPrivateIPAddressCount_updateInPlace(t *testing.T) {
+	ctx := acctest.Context(t)
+	var natGateway awstypes.NatGateway
+	resourceName := "aws_nat_gateway.test"
+	rName := acctest.RandomWithPrefix(t, acctest.ResourcePrefix)
+
+	countStart := 0
+	countUp := 3
+	countDown := 1
+
+	acctest.ParallelTest(ctx, t, resource.TestCase{
+		PreCheck:                 func() { acctest.PreCheck(ctx, t) },
+		ErrorCheck:               acctest.ErrorCheck(t, names.EC2ServiceID),
+		ProtoV5ProviderFactories: acctest.ProtoV5ProviderFactories,
+		CheckDestroy:             testAccCheckNATGatewayDestroy(ctx, t),
+		Steps: []resource.TestStep{
+			{
+				Config: testAccVPCNATGatewayConfig_secondaryPrivateIPAddressCount(rName, countStart),
+				Check: resource.ComposeAggregateTestCheckFunc(
+					testAccCheckNATGatewayExists(ctx, t, resourceName, &natGateway),
+					resource.TestCheckResourceAttr(resourceName, "secondary_private_ip_address_count", strconv.Itoa(countStart)),
+				),
+				ConfigPlanChecks: resource.ConfigPlanChecks{
+					PreApply: []plancheck.PlanCheck{
+						plancheck.ExpectResourceAction(resourceName, plancheck.ResourceActionCreate),
+					},
+				},
+			},
+			{
+				Config: testAccVPCNATGatewayConfig_secondaryPrivateIPAddressCount(rName, countUp),
+				Check: resource.ComposeAggregateTestCheckFunc(
+					testAccCheckNATGatewayExists(ctx, t, resourceName, &natGateway),
+					resource.TestCheckResourceAttr(resourceName, "secondary_private_ip_address_count", strconv.Itoa(countUp)),
+				),
+				ConfigPlanChecks: resource.ConfigPlanChecks{
+					PreApply: []plancheck.PlanCheck{
+						plancheck.ExpectResourceAction(resourceName, plancheck.ResourceActionUpdate),
+					},
+				},
+			},
+			{
+				Config: testAccVPCNATGatewayConfig_secondaryPrivateIPAddressCount(rName, countDown),
+				Check: resource.ComposeAggregateTestCheckFunc(
+					testAccCheckNATGatewayExists(ctx, t, resourceName, &natGateway),
+					resource.TestCheckResourceAttr(resourceName, "secondary_private_ip_address_count", strconv.Itoa(countDown)),
+				),
+				ConfigPlanChecks: resource.ConfigPlanChecks{
+					PreApply: []plancheck.PlanCheck{
+						plancheck.ExpectResourceAction(resourceName, plancheck.ResourceActionUpdate),
+					},
+				},
+			},
+		},
+	})
+}
+
 func TestAccVPCNATGateway_secondaryPrivateIPAddresses(t *testing.T) {
 	ctx := acctest.Context(t)
 	var natGateway awstypes.NatGateway

internal/service/ec2/wait.go

@@ -1514,6 +1514,31 @@ func waitNATGatewayAddressUnassigned(ctx context.Context, conn *ec2.Client, natG
 	return nil, err
 }
 
+func waitNATGatewaySecondaryPrivateIPAddressCount(ctx context.Context, conn *ec2.Client, natGatewayID string, expectedCount int, timeout time.Duration) (*awstypes.NatGateway, error) {
+	stateConf := &retry.StateChangeConf{
+		Pending:                   []string{"pending"},
+		Target:                    []string{"ready"},
+		Refresh:                   statusNATGatewaySecondaryPrivateIPAddressCount(conn, natGatewayID, expectedCount),
+		Timeout:                   timeout,
+		ContinuousTargetOccurence: 4,
+	}
+
+	outputRaw, err := stateConf.WaitForStateContext(ctx)
+
+	if output, ok := outputRaw.(*awstypes.NatGateway); ok {
+		for _, natGatewayAddress := range output.NatGatewayAddresses {
+			if !aws.ToBool(natGatewayAddress.IsPrimary) && natGatewayAddress.Status == awstypes.NatGatewayAddressStatusFailed {
+				retry.SetLastError(err, errors.New(aws.ToString(natGatewayAddress.FailureMessage)))
+				break
+			}
+		}
+
+		return output, err
+	}
+
+	return nil, err
+}
+
 func waitNATGatewayCreated(ctx context.Context, conn *ec2.Client, id string, timeout time.Duration) (*awstypes.NatGateway, error) {
 	stateConf := &retry.StateChangeConf{
 		Pending: enum.Slice(awstypes.NatGatewayStatePending),