From 2a4a49fd796f1541c3fc6387f05b818c1f159d14 Mon Sep 17 00:00:00 2001
From: Yury Bushmelev <jay4mail@gmail.com>
Date: Mon, 21 Aug 2023 14:14:19 +0400
Subject: [PATCH] feat!: allow to override umig availability zones

Sometimes it's profitable to specify exact list of AZs to use instead of
spreading instances across the available AZs.
Another reason is if you want to iterate over AZs subset in the specific
order (e.g. create 1st VM in a zone different from 1st available).
---
 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     = []
+}