From 4dc9d57cb81128b9f487d7f203831a8d4c6ecfd0 Mon Sep 17 00:00:00 2001
From: Yury Bushmelev <jay4mail@gmail.com>
Date: Tue, 5 Sep 2023 12:07:28 +0800
Subject: [PATCH] feat!: allow to override umig availability zones (#329)

Co-authored-by: Awais Malik <awmalik@google.com>
---
 modules/umig/README.md    |  1 +
 modules/umig/main.tf      | 10 ++++++----
 modules/umig/variables.tf |  6 ++++++
 3 files changed, 13 insertions(+), 4 deletions(-)

diff --git a/modules/umig/README.md b/modules/umig/README.md
index ac870014..7a5862a5 100644
--- a/modules/umig/README.md
+++ b/modules/umig/README.md
@@ -29,6 +29,7 @@ See the [simple](https://github.com/terraform-google-modules/terraform-google-vm
 | static\_ips | List of static IPs for VM instances | `list(string)` | `[]` | no |
 | subnetwork | Subnet to deploy to. Only one of network or subnetwork should be specified. | `string` | `""` | no |
 | subnetwork\_project | The project that subnetwork belongs to | `string` | `""` | no |
+| zones | Override the availability zones list to create the resources in. | `list(string)` | `[]` | no |
 
 ## Outputs
 
diff --git a/modules/umig/main.tf b/modules/umig/main.tf
index 12c30094..545edc91 100644
--- a/modules/umig/main.tf
+++ b/modules/umig/main.tf
@@ -23,9 +23,11 @@ locals {
   # determine type" error when var.static_ips is empty
   static_ips = concat(var.static_ips, ["NOT_AN_IP"])
 
+  zones = length(var.zones) == 0 ? data.google_compute_zones.available.names : var.zones
+
   instance_group_count = min(
     local.num_instances,
-    length(data.google_compute_zones.available.names),
+    length(local.zones),
   )
 }
 
@@ -48,7 +50,7 @@ resource "google_compute_instance_from_template" "compute_instance" {
   count    = local.num_instances
   name     = format("%s%s%s", local.hostname, var.hostname_suffix_separator, format("%03d", count.index + 1))
   project  = var.project_id
-  zone     = data.google_compute_zones.available.names[count.index % length(data.google_compute_zones.available.names)]
+  zone     = local.zones[count.index % length(local.zones)]
 
   network_interface {
     network            = var.network
@@ -105,11 +107,11 @@ resource "google_compute_instance_group" "instance_group" {
   count    = local.instance_group_count
   name     = "${local.hostname}-instance-group-${format("%03d", count.index + 1)}"
   project  = var.project_id
-  zone     = element(data.google_compute_zones.available.names, count.index)
+  zone     = element(local.zones, count.index)
   instances = matchkeys(
     google_compute_instance_from_template.compute_instance.*.self_link,
     google_compute_instance_from_template.compute_instance.*.zone,
-    [data.google_compute_zones.available.names[count.index]],
+    [local.zones[count.index]],
   )
 
   dynamic "named_port" {
diff --git a/modules/umig/variables.tf b/modules/umig/variables.tf
index a58a7b32..4c056106 100644
--- a/modules/umig/variables.tf
+++ b/modules/umig/variables.tf
@@ -115,3 +115,9 @@ variable "hostname_suffix_separator" {
   description = "Separator character to compose hostname when add_hostname_suffix is set to true."
   default     = "-"
 }
+
+variable "zones" {
+  type        = list(string)
+  description = "(Optional) List of availability zones to create VM instances in"
+  default     = []
+}