Fixes and enhancements after the bug bash. (#10223) (#7120)

* remove float from documentation

* change storage-gb type to int

* Revert "change storage-gb type to int"

This reverts commit 89795a80889287fb287497fbcbbfdbb365d40884.

* add support for dag_processor count

* fix int64 unpacking

* add validation for composer internal ip range size

* add upper limit for dag_prcessor.count

[upstream:f4abf45f154501fe04509c2da704ac0dc0d4400e]

Signed-off-by: Modular Magician <[email protected]>

Author: modular-magician

.changelog/10223.txt

@@ -0,0 +1,3 @@
+```release-note:enhancement
+composer: small fixes and enhancements after the bugbash.
+```

google-beta/services/composer/resource_composer_environment.go

@@ -6,6 +6,7 @@ import (
 	"context"
 	"fmt"
 	"log"
+	"net"
 	"regexp"
 	"strings"
 	"time"
@@ -362,11 +363,12 @@ func ResourceComposerEnvironment() *schema.Resource {
 										},
 									},
 									"composer_internal_ipv4_cidr_block": {
-										Type:        schema.TypeString,
-										Computed:    true,
-										Optional:    true,
-										ForceNew:    true,
-										Description: `IPv4 cidr range that will be used by Composer internal components.`,
+										Type:         schema.TypeString,
+										Computed:     true,
+										Optional:     true,
+										ForceNew:     true,
+										ValidateFunc: validateComposerInternalIpv4CidrBlock,
+										Description:  `IPv4 cidr range that will be used by Composer internal components.`,
 									},
 								},
 							},
@@ -936,6 +938,14 @@ func ResourceComposerEnvironment() *schema.Resource {
 													ValidateFunc: validation.FloatAtLeast(0),
 													Description:  `Storage (GB) request and limit for DAG processor.`,
 												},
+												"count": {
+													Type:         schema.TypeInt,
+													Optional:     true,
+													ForceNew:     false,
+													Computed:     true,
+													ValidateFunc: validation.IntBetween(0, 3),
+													Description:  `Number of DAG processors.`,
+												},
 											},
 										},
 									},
@@ -1838,6 +1848,7 @@ func flattenComposerEnvironmentConfigWorkloadsConfig(workloadsConfig *composer.W
 		transformedDagProcessor["cpu"] = wlCfgDagProcessor.Cpu
 		transformedDagProcessor["memory_gb"] = wlCfgDagProcessor.MemoryGb
 		transformedDagProcessor["storage_gb"] = wlCfgDagProcessor.StorageGb
+		transformedDagProcessor["count"] = wlCfgDagProcessor.Count
 	}
 
 	transformed["scheduler"] = []interface{}{transformedScheduler}
@@ -2330,6 +2341,7 @@ func expandComposerEnvironmentConfigWorkloadsConfig(v interface{}, d *schema.Res
 			transformedDagProcessor.Cpu = originalDagProcessorRaw["cpu"].(float64)
 			transformedDagProcessor.MemoryGb = originalDagProcessorRaw["memory_gb"].(float64)
 			transformedDagProcessor.StorageGb = originalDagProcessorRaw["storage_gb"].(float64)
+			transformedDagProcessor.Count = int64(originalDagProcessorRaw["count"].(int))
 			transformed.DagProcessor = transformedDagProcessor
 		}
 	}
@@ -3018,3 +3030,17 @@ func versionValidationCustomizeDiffFunc(ctx context.Context, d *schema.ResourceD
 	}
 	return nil
 }
+
+func validateComposerInternalIpv4CidrBlock(v any, k string) (warns []string, errs []error) {
+	cidr_range := v.(string)
+	_, ip_net, err := net.ParseCIDR(cidr_range)
+	if err != nil {
+		errs = append(errs, fmt.Errorf("Invalid CIDR range: %s", err))
+		return
+	}
+	ones, _ := ip_net.Mask.Size()
+	if ones != 20 {
+		errs = append(errs, fmt.Errorf("Composer Internal IPv4 CIDR range must have size /20"))
+	}
+	return
+}

google-beta/services/composer/resource_composer_environment_test.go

@@ -3141,6 +3141,7 @@ resource "google_composer_environment" "test" {
         cpu          = 1
         memory_gb    = 2.5
         storage_gb   = 2
+        count        = 1
       }
     }
     enable_private_environment = true
@@ -3184,6 +3185,7 @@ resource "google_composer_environment" "test" {
         cpu          = 2
         memory_gb    = 2
         storage_gb   = 1
+        count        = 2
       }
     }
     enable_private_environment = false

website/docs/r/composer_environment.html.markdown

@@ -1371,7 +1371,7 @@ The `worker` block supports:
   (Optional)
   The amount of memory (GB) for a single Airflow worker.
 
-* `float storage_gb`
+* `storage_gb`
   (Optional)
   The amount of storage (GB) for a single Airflow worker.
 
@@ -1396,7 +1396,7 @@ The `dag_processor` block supports:
   (Optional)
   Memory (GB) request and limit for DAG processor.
 
-* `float storage_gb`
+* `storage_gb`
   (Optional)
   Storage (GB) request and limit for DAG processor.