From 997c48b1de197ca5c952b5660adf17ec41991241 Mon Sep 17 00:00:00 2001
From: Daniel Panzella <daniel.panzella@wandb.com>
Date: Mon, 30 Sep 2024 18:58:59 -0700
Subject: [PATCH] fix: Handle per subscription AZ restrictions (#106)

---
 examples/public-dns/outputs.tf | 2 +-
 main.tf                        | 8 ++++++--
 outputs.tf                     | 1 +
 3 files changed, 8 insertions(+), 3 deletions(-)

diff --git a/examples/public-dns/outputs.tf b/examples/public-dns/outputs.tf
index 2cc18f1..fd8e98e 100644
--- a/examples/public-dns/outputs.tf
+++ b/examples/public-dns/outputs.tf
@@ -8,4 +8,4 @@ output "url" {
 
 output "oidc_issuer_url" {
   value = module.wandb.oidc_issuer_url
-}
\ No newline at end of file
+}
diff --git a/main.tf b/main.tf
index b8269c7..943aebd 100644
--- a/main.tf
+++ b/main.tf
@@ -127,8 +127,12 @@ locals {
       sku.name == local.kubernetes_instance_type
     )
   ]
-  num_zones       = var.node_pool_zones != null ? length(var.node_pool_zones) : var.node_pool_num_zones
-  node_pool_zones = var.node_pool_zones != null ? var.node_pool_zones : slice(sort(local.vm_skus[0].locationInfo[0].zones), 0, local.num_zones)
+
+  num_zones        = var.node_pool_zones != null ? length(var.node_pool_zones) : var.node_pool_num_zones
+  restricted_zones = length(local.vm_skus[0].restrictions) > 0 ? local.vm_skus[0].restrictions[0].restrictionInfo.zones : []
+  all_zones        = local.vm_skus[0].locationInfo[0].zones
+  valid_zones      = setsubtract(local.all_zones, local.restricted_zones)
+  node_pool_zones  = var.node_pool_zones != null ? var.node_pool_zones : slice(sort(local.valid_zones), 0, local.num_zones)
 }
 
 module "app_aks" {
diff --git a/outputs.tf b/outputs.tf
index 3176fb0..f1a9ab7 100644
--- a/outputs.tf
+++ b/outputs.tf
@@ -64,3 +64,4 @@ output "client_id" {
 output "tenant_id" {
   value = module.identity.identity.tenant_id
 }
+