From 6d765866a99a07e3cb0bf0ccc9eaa73612b2f3c4 Mon Sep 17 00:00:00 2001
From: Lei Jin <166442440+leijin-lw@users.noreply.github.com>
Date: Mon, 3 Jun 2024 09:10:29 -0700
Subject: [PATCH] feat(GROW-2931): Add support for GCP default labels (#1633)

Signed-off-by: Lei Jin <lei.jin@lacework.net>
---
 lwgenerate/gcp/gcp.go      | 32 +++++++++++++++++++++++++++++---
 lwgenerate/gcp/gcp_test.go | 26 ++++++++++++++++++++++++++
 lwgenerate/gcp/gke.go      |  3 +++
 3 files changed, 58 insertions(+), 3 deletions(-)

diff --git a/lwgenerate/gcp/gcp.go b/lwgenerate/gcp/gcp.go
index 32669d8fe..924238f06 100644
--- a/lwgenerate/gcp/gcp.go
+++ b/lwgenerate/gcp/gcp.go
@@ -141,6 +141,9 @@ type GenerateGcpTfConfigurationArgs struct {
 
 	Projects []string
 
+	// Default GCP Provider labels
+	ProviderDefaultLabels map[string]interface{}
+
 	// Add custom blocks to the root `terraform{}` block. Can be used for advanced configuration. Things like backend, etc
 	ExtraBlocksRootTerraform []*hclwrite.Block
 
@@ -235,6 +238,13 @@ func WithGcpServiceAccountCredentials(path string) GcpTerraformModifier {
 	}
 }
 
+// WithProviderDefaultLabels adds default_labels to the provider configuration for GCP (if labels are present)
+func WithProviderDefaultLabels(labels map[string]interface{}) GcpTerraformModifier {
+	return func(c *GenerateGcpTfConfigurationArgs) {
+		c.ProviderDefaultLabels = labels
+	}
+}
+
 // WithConfigOutputs Set Custom Terraform Outputs
 func WithCustomOutputs(outputs []lwgenerate.HclOutput) GcpTerraformModifier {
 	return func(c *GenerateGcpTfConfigurationArgs) {
@@ -465,7 +475,7 @@ func (args *GenerateGcpTfConfigurationArgs) Generate() (string, error) {
 	}
 
 	gcpProvider, err := createGcpProvider(args.ExtraProviderArguments,
-		args.ServiceAccountCredentials, args.GcpProjectId, args.Regions, "")
+		args.ServiceAccountCredentials, args.GcpProjectId, args.Regions, "", args.ProviderDefaultLabels)
 	if err != nil {
 		return "", errors.Wrap(err, "failed to generate gcp provider")
 	}
@@ -546,6 +556,7 @@ func createGcpProvider(
 	projectId string,
 	regionsArg []string,
 	alias string,
+	providerDefaultLabels map[string]interface{},
 ) ([]*hclwrite.Block, error) {
 	blocks := []*hclwrite.Block{}
 
@@ -578,10 +589,25 @@ func createGcpProvider(
 			attrs["region"] = region
 		}
 
+		modifiers := []lwgenerate.HclProviderModifier{
+			lwgenerate.HclProviderWithAttributes(attrs),
+		}
+
+		if len(providerDefaultLabels) != 0 {
+			defaultLabelsBlock, err := lwgenerate.HclCreateGenericBlock(
+				"default_labels",
+				nil,
+				providerDefaultLabels,
+			)
+			if err != nil {
+				return nil, err
+			}
+			modifiers = append(modifiers, lwgenerate.HclProviderWithGenericBlocks(defaultLabelsBlock))
+		}
+
 		provider, err := lwgenerate.NewProvider(
 			"google",
-			lwgenerate.HclProviderWithAttributes(attrs),
-		).ToBlock()
+			modifiers...).ToBlock()
 		if err != nil {
 			return nil, err
 		}
diff --git a/lwgenerate/gcp/gcp_test.go b/lwgenerate/gcp/gcp_test.go
index 977277c3c..82f0ad434 100644
--- a/lwgenerate/gcp/gcp_test.go
+++ b/lwgenerate/gcp/gcp_test.go
@@ -862,6 +862,19 @@ func TestGenerationConfigWithOutputs(t *testing.T) {
 	assert.Equal(t, RequiredProviders+"\n"+gcpProvider+"\n"+moduleImportProjectLevelAuditLogWithoutConfiguration+"\n"+customOutput, hcl)
 }
 
+func TestGenerationConfigWithDefaultProviderLabels(t *testing.T) {
+	hcl, err := gcp.NewTerraform(
+		false, false, true, false,
+		gcp.WithGcpServiceAccountCredentials("/path/to/credentials"),
+		gcp.WithProjectId(projectName),
+		gcp.WithRegions([]string{"us-east1"}),
+		gcp.WithProviderDefaultLabels(map[string]interface{}{"LABEL_TEST": "foo", "LABEL_TEST1": "bar"})).Generate()
+	assert.Nil(t, err)
+	assert.NotNil(t, hcl)
+	assert.Equal(t, RequiredProviders+"\n"+gcpProviderWithDefaultLabels+"\n"+moduleImportProjectLevelAuditLogWithoutConfiguration, hcl)
+
+}
+
 func ProviderWithCredentials(projectName string) string {
 	return fmt.Sprintf(`provider "google" {
   credentials = "/path/to/credentials"
@@ -905,6 +918,19 @@ var gcpProviderWithExtraArguments = `provider "google" {
 }
 `
 
+var gcpProviderWithDefaultLabels = `provider "google" {
+  alias       = "us-east1"
+  credentials = "/path/to/credentials"
+  project     = "project1"
+  region      = "us-east1"
+
+  default_labels {
+    LABEL_TEST  = "foo"
+    LABEL_TEST1 = "bar"
+  }
+}
+`
+
 var testVariable = `variable "var_name" {
 }
 `
diff --git a/lwgenerate/gcp/gke.go b/lwgenerate/gcp/gke.go
index fb7cce228..0413b3ca9 100644
--- a/lwgenerate/gcp/gke.go
+++ b/lwgenerate/gcp/gke.go
@@ -24,6 +24,8 @@ type GenerateGkeTfConfigurationArgs struct {
 	PubSubTopicLabels            map[string]string
 	ServiceAccountCredentials    string
 	WaitTime                     string
+	// Default GCP Provider labels
+	ProviderDefaultLabels map[string]interface{}
 	// Add custom blocks to the root `terraform{}` block. Can be used for advanced configuration. Things like backend, etc
 	ExtraBlocksRootTerraform []*hclwrite.Block
 	// ExtraProviderArguments allows adding more arguments to the provider block as needed (custom use cases)
@@ -48,6 +50,7 @@ func (args *GenerateGkeTfConfigurationArgs) Generate() (string, error) {
 		args.ProjectId,
 		[]string{},
 		args.GcpProviderAlias,
+		args.ProviderDefaultLabels,
 	)
 	if err != nil {
 		return "", errors.Wrap(err, "failed to generate gcp provider")