fix: allow specifying number of AZ to use when autoselecting zones

main.tf

@@ -108,7 +108,7 @@ module "app_lb" {
 }
 
 data "external" "az_zones" {
-  program = ["bash", "${path.module}/vmtype_to_az.sh", local.kubernetes_instance_type, azurerm_resource_group.default.location]
+  program = ["bash", "${path.module}/vmtype_to_az.sh", local.kubernetes_instance_type, azurerm_resource_group.default.location, var.node_pool_num_zones]
 }
 
 locals {

variables.tf

@@ -218,6 +218,12 @@ variable "node_pool_zones" {
   default     = null
 }
 
+variable "node_pool_num_zones" {
+  type        = number
+  description = "Number of availability zones to use for the node pool when node_pool_zones is not set. If neither are set, 3 zones will be used"
+  default     = null
+}
+
 variable "node_max_pods" {
   type        = number
   description = "Maximum number of pods per node"

vmtype_to_az.sh

@@ -11,16 +11,22 @@
 # Copy script arguments to named environment variables
 VM_TYPE="$1"
 REGION="$2"
+NUM_ZONES="$3"
 
 # Check if both arguments are provided
 if [ -z "$VM_TYPE" ] || [ -z "$REGION" ]; then
     echo "Error: Both VM type and region must be provided." >&2
-    echo "Usage: $0 <vm_type> <region>" >&2
+    echo "Usage: $0 <vm_type> <region> <num_zones>" >&2
     exit 1
 fi
 
+# Default to 3 zones if not specified
+if [ -z "$NUM_ZONES" ]; then
+    NUM_ZONES=3
+fi
+
 # Query Azure CLI for availability zones in the region for the specified VM type
-ZONES=$(az vm list-skus --location "$REGION" --size "$VM_TYPE" --query "[0].locationInfo[0].zones" -o json | jq -r -c 'sort | .[0:3]')
+ZONES=$(az vm list-skus --location "$REGION" --size "$VM_TYPE" --query "[0].locationInfo[0].zones" -o json | jq -r -c "sort | .[0:$NUM_ZONES]")
 
 # Check if the query returned any results
 if [ -z "$ZONES" ] || [ "$ZONES" == "null" ]; then