fix: lb and lb pool config migration (#6170)

Author: vaishakdinesh

cmd/migrate/load_balancer.go

@@ -2,6 +2,7 @@ package main
 
 import (
 	"fmt"
+	"regexp"
 	"strings"
 	"github.com/cloudflare/terraform-provider-cloudflare/cmd/migrate/ast"
 	"github.com/hashicorp/hcl/v2"
@@ -32,7 +33,12 @@ func transformLoadBalancerBlock(block *hclwrite.Block, diags ast.Diagnostics) {
 	transformPoolBlocksToMap(block, "country_pools", "country")
 	transformPoolBlocksToMap(block, "pop_pools", "pop")
 
-	// Transform rules attribute to ensure region_pools.region is a list
+	// Transform dynamic rules blocks to for expressions
+	// v4: dynamic "rules" { for_each = ...; content { ... } }
+	// v5: rules = [for ... : { ... }]
+	transformDynamicRulesBlocksToAttribute(block, diags)
+	
+	// Transform rules attribute to ensure region_pools is a map
 	transformLoadBalancerRules(block, diags)
 }
 
@@ -187,77 +193,110 @@ func transformPoolArrayToMap(block *hclwrite.Block, blockName string, keyField s
 }
 
 
-// transformLoadBalancerRules transforms the rules attribute to ensure region_pools.region is a list
-func transformLoadBalancerRules(block *hclwrite.Block, diags ast.Diagnostics) {
+// transformLoadBalancerRulesString applies string-level transformations to consolidate region_pools
+func transformLoadBalancerRulesString(block *hclwrite.Block) {
 	// Get the rules attribute
 	rulesAttr := block.Body().GetAttribute("rules")
 	if rulesAttr == nil {
 		return
 	}
 
-	// Parse the rules expression
-	rulesExpr := ast.WriteExpr2Expr(rulesAttr.Expr(), diags)
+	// Get the tokens as a string
+	tokens := rulesAttr.Expr().BuildTokens(nil)
+	content := string(tokens.Bytes())
 
-	// Check if it's a tuple (list of rules)
-	tup, ok := rulesExpr.(*hclsyntax.TupleConsExpr)
-	if !ok {
-		// Can't parse - keep as is since we don't know the structure
+	// Check if we have region_pools to process
+	if !strings.Contains(content, "region_pools") {
 		return
 	}
 
-	// Process each rule in the list
-	modified := false
-	for _, ruleExpr := range tup.Exprs {
-		// Check if the rule is an object
-		ruleObj, ok := ruleExpr.(*hclsyntax.ObjectConsExpr)
-		if !ok {
+	// Process each overrides block separately to avoid mixing rules
+	// Find all overrides blocks
+	overridesPattern := `overrides\s*=\s*\{[^{}]*(?:\{[^}]*\}[^{}]*)*\}`
+	overridesMatches := regexp.MustCompile(overridesPattern).FindAllString(content, -1)
+	
+	for _, overridesBlock := range overridesMatches {
+		// Check if this overrides block has region_pools
+		if !strings.Contains(overridesBlock, "region_pools") {
 			continue
 		}
 
-		// Find the overrides attribute
-		for _, item := range ruleObj.Items {
-			keyExpr, ok := item.KeyExpr.(*hclsyntax.ObjectConsKeyExpr)
-			if !ok {
-				continue
-			}
+		// Find all region_pools within this specific overrides block
+		regionPoolsPattern := `region_pools\s*=\s*\{[^}]*?\}`
+		regionPoolsMatches := regexp.MustCompile(regionPoolsPattern).FindAllString(overridesBlock, -1)
+		
+		if len(regionPoolsMatches) == 0 {
+			continue
+		}
+		
+		// Parse each region_pools to extract region and pool_ids
+		var regionPoolData [][]string
+		for _, rpMatch := range regionPoolsMatches {
+			// Extract region
+			regionPattern := `region\s*=\s*"([^"]+)"`
+			regionMatch := regexp.MustCompile(regionPattern).FindStringSubmatch(rpMatch)
 
-			if ast.Expr2S(keyExpr, diags) == "overrides" {
-				// Process the overrides object
-				overridesObj, ok := item.ValueExpr.(*hclsyntax.ObjectConsExpr)
-				if !ok {
-					continue
-				}
-				
-				// Find region_pools in overrides
-				for _, overrideItem := range overridesObj.Items {
-					overrideKeyExpr, ok := overrideItem.KeyExpr.(*hclsyntax.ObjectConsKeyExpr)
-					if !ok {
-						continue
-					}
-					
-					if ast.Expr2S(overrideKeyExpr, diags) == "region_pools" {
-						// Transform region_pools
-						if transformRegionPools(&overrideItem.ValueExpr, diags) {
-							modified = true
-						}
-					}
+			// Extract pool_ids
+			poolIDsPattern := `pool_ids\s*=\s*(\[(?:[^\[\]]*|\[[^\]]*\])*\])`
+			poolIDsMatch := regexp.MustCompile(poolIDsPattern).FindStringSubmatch(rpMatch)
+			
+			if regionMatch != nil && poolIDsMatch != nil {
+				regionPoolData = append(regionPoolData, []string{rpMatch, regionMatch[1], poolIDsMatch[1]})
+			}
+		}
+		
+		// If we found region_pools to consolidate
+		if len(regionPoolData) > 0 {
+			// Build the consolidated map
+			consolidatedMap := "region_pools = {\n"
+			for i, data := range regionPoolData {
+				region := data[1]
+				poolIds := data[2]
+				if i < len(regionPoolData)-1 {
+					consolidatedMap += fmt.Sprintf("      %s = %s,\n", region, poolIds)
+				} else {
+					consolidatedMap += fmt.Sprintf("      %s = %s\n", region, poolIds)
 				}
 			}
+			consolidatedMap += "    }"
+			
+			// Create modified overrides block
+			modifiedOverrides := overridesBlock
+			
+			// Replace first region_pools with consolidated map
+			modifiedOverrides = strings.Replace(modifiedOverrides, regionPoolData[0][0], consolidatedMap, 1)
+			
+			// Remove all other region_pools blocks
+			for i := 1; i < len(regionPoolData); i++ {
+				modifiedOverrides = strings.Replace(modifiedOverrides, regionPoolData[i][0], "", 1)
+			}
+			
+			// Clean up extra whitespace and commas
+			modifiedOverrides = regexp.MustCompile(`\n\s*\n\s*\n`).ReplaceAllString(modifiedOverrides, "\n")
+			
+			// Replace the original overrides block with the modified one
+			content = strings.Replace(content, overridesBlock, modifiedOverrides, 1)
 		}
 	}
 
-	// Only update if we made changes
-	if modified {
-		// Convert the modified expression back to a string and reparse as HCL
-		exprStr := ast.Expr2S(rulesExpr, diags)
-		// Use hclwrite to parse and format the expression properly
-		newExpr, err := hclwrite.ParseConfig([]byte("temp = "+exprStr), "temp.hcl", hcl.InitialPos)
-		if err == nil && newExpr.Body().Attributes()["temp"] != nil {
-			block.Body().SetAttributeRaw("rules", newExpr.Body().Attributes()["temp"].Expr().BuildTokens(nil))
-		}
+	// Parse the modified content and set it back
+	tempConfig := fmt.Sprintf("temp = %s", content)
+	tempFile, err := hclwrite.ParseConfig([]byte(tempConfig), "temp.hcl", hcl.InitialPos)
+	if err == nil && tempFile.Body().GetAttribute("temp") != nil {
+		block.Body().SetAttributeRaw("rules", tempFile.Body().GetAttribute("temp").Expr().BuildTokens(nil))
 	}
 }
 
+
+// transformLoadBalancerRules transforms the rules attribute to:
+// 1. Consolidate multiple region_pools objects into a single map
+// 2. Fix the structure from { region = "X", pool_ids = [...] } to { "X" = [...] }
+func transformLoadBalancerRules(block *hclwrite.Block, diags ast.Diagnostics) {
+	// Just use the string-level transformation which handles the region_pools consolidation
+	// This preserves complex expressions that can't be parsed
+	transformLoadBalancerRulesString(block)
+}
+
 // transformRegionPools ensures that each region_pools element has region as a list
 func transformRegionPools(expr *hclsyntax.Expression, diags ast.Diagnostics) bool {
 	modified := false
@@ -314,3 +353,201 @@ func transformSingleRegionPool(poolExpr hclsyntax.Expression, diags ast.Diagnost
 	return modified
 }
 
+// extractRegionPoolIntoMap extracts region and pool_ids from a region_pool object and adds to map
+func extractRegionPoolIntoMap(poolObj *hclsyntax.ObjectConsExpr, regionPoolsMap map[string]hclsyntax.Expression, diags ast.Diagnostics) {
+	var regionValues []string
+	var poolIDsExpr hclsyntax.Expression
+	
+	for _, item := range poolObj.Items {
+		keyExpr, ok := item.KeyExpr.(*hclsyntax.ObjectConsKeyExpr)
+		if !ok {
+			continue
+		}
+		
+		keyName := ast.Expr2S(keyExpr, diags)
+		switch keyName {
+		case "region":
+			// Extract region value(s)
+			
+			// Convert the expression to string regardless of type
+			regionStr := ast.Expr2S(item.ValueExpr, diags)
+			
+			// Clean up the value - remove quotes and whitespace
+			regionStr = strings.TrimSpace(regionStr)
+			if len(regionStr) >= 2 && regionStr[0] == '"' && regionStr[len(regionStr)-1] == '"' {
+				regionStr = regionStr[1 : len(regionStr)-1]
+			}
+			
+			if regionStr != "" {
+				regionValues = append(regionValues, regionStr)
+			}
+		case "pool_ids":
+			poolIDsExpr = item.ValueExpr
+		}
+	}
+	
+	// Add to map for each region
+	for _, region := range regionValues {
+		if poolIDsExpr != nil {
+			regionPoolsMap[region] = poolIDsExpr
+		}
+	}
+}
+
+// extractPoolIntoMap extracts key field and pool_ids from a pool object and adds to map
+func extractPoolIntoMap(poolObj *hclsyntax.ObjectConsExpr, keyField string, poolMap map[string]hclsyntax.Expression, diags ast.Diagnostics) {
+	var keyValue string
+	var poolIDsExpr hclsyntax.Expression
+	
+	for _, item := range poolObj.Items {
+		keyExpr, ok := item.KeyExpr.(*hclsyntax.ObjectConsKeyExpr)
+		if !ok {
+			continue
+		}
+		
+		keyName := ast.Expr2S(keyExpr, diags)
+		if keyName == keyField {
+			// Extract key value
+			if litExpr, ok := item.ValueExpr.(*hclsyntax.LiteralValueExpr); ok {
+				if litExpr.Val.Type().FriendlyName() == "string" {
+					keyValue = litExpr.Val.AsString()
+				}
+			}
+		} else if keyName == "pool_ids" {
+			poolIDsExpr = item.ValueExpr
+		}
+	}
+	
+	// Add to map
+	if keyValue != "" && poolIDsExpr != nil {
+		poolMap[keyValue] = poolIDsExpr
+	}
+}
+
+// transformDynamicRulesBlocksToAttribute converts dynamic "rules" blocks to a rules attribute with for expression
+func transformDynamicRulesBlocksToAttribute(block *hclwrite.Block, diags ast.Diagnostics) {
+	body := block.Body()
+	dynamicRulesBlocks := []*hclwrite.Block{}
+	
+	// Find all dynamic blocks with label "rules"
+	for _, childBlock := range body.Blocks() {
+		if childBlock.Type() == "dynamic" && len(childBlock.Labels()) > 0 && childBlock.Labels()[0] == "rules" {
+			dynamicRulesBlocks = append(dynamicRulesBlocks, childBlock)
+		}
+	}
+	
+	if len(dynamicRulesBlocks) == 0 {
+		return
+	}
+	
+	// Process each dynamic rules block
+	for _, dynBlock := range dynamicRulesBlocks {
+		dynBody := dynBlock.Body()
+		
+		// Get the for_each expression
+		forEachAttr := dynBody.GetAttribute("for_each")
+		if forEachAttr == nil {
+			continue
+		}
+		
+		// Find the content block
+		var contentBlock *hclwrite.Block
+		for _, cb := range dynBody.Blocks() {
+			if cb.Type() == "content" {
+				contentBlock = cb
+				break
+			}
+		}
+		
+		if contentBlock == nil {
+			continue
+		}
+		
+		// Get the for_each expression tokens
+		forEachTokens := forEachAttr.Expr().BuildTokens(nil)
+		
+		// Build the for expression
+		// rules = [for rule in <for_each> : { <content> }]
+		tokens := hclwrite.Tokens{
+			&hclwrite.Token{Type: hclsyntax.TokenOBrack, Bytes: []byte("[")},
+			&hclwrite.Token{Type: hclsyntax.TokenIdent, Bytes: []byte("for")},
+			&hclwrite.Token{Type: hclsyntax.TokenIdent, Bytes: []byte(" rules")},
+			&hclwrite.Token{Type: hclsyntax.TokenIdent, Bytes: []byte(" in")},
+			&hclwrite.Token{Type: hclsyntax.TokenIdent, Bytes: []byte(" ")},
+		}
+		
+		// Add the for_each expression
+		tokens = append(tokens, forEachTokens...)
+		tokens = append(tokens,
+			&hclwrite.Token{Type: hclsyntax.TokenIdent, Bytes: []byte(" ")},
+			&hclwrite.Token{Type: hclsyntax.TokenColon, Bytes: []byte(":")},
+			&hclwrite.Token{Type: hclsyntax.TokenIdent, Bytes: []byte(" ")},
+			&hclwrite.Token{Type: hclsyntax.TokenOBrace, Bytes: []byte("{")},
+			&hclwrite.Token{Type: hclsyntax.TokenNewline, Bytes: []byte("\n")},
+		)
+		
+		// Add content block attributes as object properties
+		contentBody := contentBlock.Body()
+		
+		// Process attributes in the content block
+		attrs := contentBody.Attributes()
+		first := true
+		for name, attr := range attrs {
+			if !first {
+				tokens = append(tokens, &hclwrite.Token{Type: hclsyntax.TokenComma, Bytes: []byte(",")})
+				tokens = append(tokens, &hclwrite.Token{Type: hclsyntax.TokenNewline, Bytes: []byte("\n")})
+			}
+			first = false
+			
+			// Add indentation
+			tokens = append(tokens, &hclwrite.Token{Type: hclsyntax.TokenIdent, Bytes: []byte("      ")})
+			
+			// Add attribute name
+			tokens = append(tokens, &hclwrite.Token{Type: hclsyntax.TokenIdent, Bytes: []byte(name)})
+			tokens = append(tokens, &hclwrite.Token{Type: hclsyntax.TokenEqual, Bytes: []byte(" = ")})
+			
+			// Get the attribute value expression tokens
+			attrTokens := attr.Expr().BuildTokens(nil)
+			
+			// For rules, we use a simpler replacement since "rules" is used as the entire iterator
+			// not "rules.key" and "rules.value" like in the origins case
+			// We just replace "rules.value" -> "rules" and "rules.key" -> "rules"
+			attrStr := string(attrTokens.Bytes())
+			// Only replace when preceded by a dot to avoid replacing other uses of "rules"
+			attrStr = strings.ReplaceAll(attrStr, ".rules.value", ".rules")
+			attrStr = strings.ReplaceAll(attrStr, ".rules.key", ".rules")
+			// Also handle at the start of expression
+			if strings.HasPrefix(attrStr, "rules.value") {
+				attrStr = "rules" + attrStr[11:] // len("rules.value") = 11
+			}
+			if strings.HasPrefix(attrStr, "rules.key") {
+				attrStr = "rules" + attrStr[9:] // len("rules.key") = 9
+			}
+			
+			// Parse the modified expression
+			tempConfig := fmt.Sprintf("temp = %s", attrStr)
+			tempFile, parseErr := hclwrite.ParseConfig([]byte(tempConfig), "temp.hcl", hcl.InitialPos)
+			if parseErr == nil && tempFile.Body().GetAttribute("temp") != nil {
+				tokens = append(tokens, tempFile.Body().GetAttribute("temp").Expr().BuildTokens(nil)...)
+			} else {
+				// Fall back to original tokens if parsing fails
+				tokens = append(tokens, attrTokens...)
+			}
+		}
+		
+		// Close the object and list
+		tokens = append(tokens,
+			&hclwrite.Token{Type: hclsyntax.TokenNewline, Bytes: []byte("\n")},
+			&hclwrite.Token{Type: hclsyntax.TokenIdent, Bytes: []byte("    ")},
+			&hclwrite.Token{Type: hclsyntax.TokenCBrace, Bytes: []byte("}")},
+			&hclwrite.Token{Type: hclsyntax.TokenCBrack, Bytes: []byte("]")},
+		)
+		
+		// Remove the dynamic block
+		body.RemoveBlock(dynBlock)
+		
+		// Add the rules attribute
+		body.SetAttributeRaw("rules", tokens)
+	}
+}
+