From 6d9c0e68939719a846485f8e98fb94ad284cc08d Mon Sep 17 00:00:00 2001
From: Daniel Panzella <daniel.panzella@wandb.com>
Date: Mon, 16 Sep 2024 15:01:57 -0700
Subject: [PATCH] fix: allow specifying number of AZ to use when autoselecting
 zones

---
 main.tf         |  2 +-
 variables.tf    |  6 ++++++
 vmtype_to_az.sh | 10 ++++++++--
 3 files changed, 15 insertions(+), 3 deletions(-)

diff --git a/main.tf b/main.tf
index 1f6c3f5..febd734 100644
--- a/main.tf
+++ b/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 {
diff --git a/variables.tf b/variables.tf
index 8c9c7c9..c543998 100644
--- a/variables.tf
+++ b/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"
diff --git a/vmtype_to_az.sh b/vmtype_to_az.sh
index b69a196..ae74523 100755
--- a/vmtype_to_az.sh
+++ b/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