From ae4c2bdad72bd97f0ce67f59c6e7b61f4ed80bc3 Mon Sep 17 00:00:00 2001
From: Thomas Rodgers <thomasrodgers@google.com>
Date: Mon, 24 Jun 2024 13:40:16 -0700
Subject: [PATCH] Make region from zone functions support more zone names
 (#10714)

---
 .../terraform/functions/region_from_zone.go    | 18 ++++++++++--------
 .../third_party/terraform/tpgresource/utils.go | 11 +++++------
 .../tests/data/example_pubsub_lite_topic.tf    |  2 +-
 3 files changed, 16 insertions(+), 15 deletions(-)

diff --git a/mmv1/third_party/terraform/functions/region_from_zone.go b/mmv1/third_party/terraform/functions/region_from_zone.go
index d1ba3104e88c..1a75b751169b 100644
--- a/mmv1/third_party/terraform/functions/region_from_zone.go
+++ b/mmv1/third_party/terraform/functions/region_from_zone.go
@@ -3,6 +3,7 @@ package functions
 import (
 	"context"
 	"fmt"
+	"strings"
 
 	"github.com/hashicorp/terraform-plugin-framework/function"
 )
@@ -35,23 +36,24 @@ func (f RegionFromZoneFunction) Definition(ctx context.Context, req function.Def
 
 func (f RegionFromZoneFunction) Run(ctx context.Context, req function.RunRequest, resp *function.RunResponse) {
 	// Load arguments from function call
-	var arg0 string
-	resp.Error = function.ConcatFuncErrors(req.Arguments.GetArgument(ctx, 0, &arg0))
+	var zone string
+	resp.Error = function.ConcatFuncErrors(req.Arguments.GetArgument(ctx, 0, &zone))
 	if resp.Error != nil {
 		return
 	}
 
-	if arg0 == "" {
+	if zone == "" {
 		err := function.NewArgumentFuncError(0, "The input string cannot be empty.")
 		resp.Error = function.ConcatFuncErrors(err)
 		return
 	}
 
-	if arg0[len(arg0)-2] != '-' {
-		err := function.NewArgumentFuncError(0, fmt.Sprintf("The input string \"%s\" is not a valid zone name.", arg0))
+	zoneParts := strings.Split(zone, "-")
+
+	if len(zoneParts) < 3 {
+		err := function.NewArgumentFuncError(0, fmt.Sprintf("The input string \"%s\" is not a valid zone name.", zone))
 		resp.Error = function.ConcatFuncErrors(err)
-		return
+	} else {
+		resp.Error = function.ConcatFuncErrors(resp.Result.Set(ctx, strings.Join(zoneParts[:len(zoneParts)-1], "-")))
 	}
-
-	resp.Error = function.ConcatFuncErrors(resp.Result.Set(ctx, arg0[:len(arg0)-2]))
 }
diff --git a/mmv1/third_party/terraform/tpgresource/utils.go b/mmv1/third_party/terraform/tpgresource/utils.go
index 4d16bc9a15eb..a86a88d98a0e 100644
--- a/mmv1/third_party/terraform/tpgresource/utils.go
+++ b/mmv1/third_party/terraform/tpgresource/utils.go
@@ -58,15 +58,14 @@ type TerraformResourceDiff interface {
 // Contains functions that don't really belong anywhere else.
 
 // GetRegionFromZone returns the region from a zone for Google cloud.
-// This is by removing the last two chars from the zone name to leave the region
-// If there aren't enough characters in the input string, an empty string is returned
+// This is by removing the characters after the last '-'.
 // e.g. southamerica-west1-a => southamerica-west1
 func GetRegionFromZone(zone string) string {
-	if zone != "" && len(zone) > 2 {
-		region := zone[:len(zone)-2]
-		return region
+	zoneParts := strings.Split(zone, "-")
+	if len(zoneParts) < 3 {
+		return ""
 	}
-	return ""
+	return strings.Join(zoneParts[:len(zoneParts)-1], "-")
 }
 
 // Infers the region based on the following (in order of priority):
diff --git a/mmv1/third_party/tgc/tests/data/example_pubsub_lite_topic.tf b/mmv1/third_party/tgc/tests/data/example_pubsub_lite_topic.tf
index 6697b15713e8..47607a535ce8 100644
--- a/mmv1/third_party/tgc/tests/data/example_pubsub_lite_topic.tf
+++ b/mmv1/third_party/tgc/tests/data/example_pubsub_lite_topic.tf
@@ -29,7 +29,7 @@ provider "google" {
 
 resource "google_pubsub_lite_topic" "example" {
   name = "example-topic"
-  zone = "us-central1a"
+  zone = "us-central1-a"
 
   partition_config {
     count = 1