Repository: ariga/atlas
Pull request #2423 (2024-01-07 21:47:50)

schemahcl: allow using references in the for_each meta argument


Patch:
diff --git a/schemahcl/context.go b/schemahcl/context.go
index 22db75ca..a90d1e84 100644
--- a/schemahcl/context.go
+++ b/schemahcl/context.go
@@ -507,7 +507,10 @@ func extractDef(blk *hclsyntax.Block, parent *blockDef) *blockDef {
 		children: make(map[string]*blockDef),
 	}
 	for _, a := range blk.Body.Attributes {
-		cur.fields[a.Name] = struct{}{}
+		// Skip meta arguments.
+		if a.Name != forEachAttr {
+			cur.fields[a.Name] = struct{}{}
+		}
 	}
 	for _, c := range blk.Body.Blocks {
 		cur.child(extractDef(c, cur))
diff --git a/schemahcl/schemahcl.go b/schemahcl/schemahcl.go
index 15e8716d..b6cd4446 100644
--- a/schemahcl/schemahcl.go
+++ b/schemahcl/schemahcl.go
@@ -294,14 +294,20 @@ func (s *State) EvalFiles(paths []string, v any, input map[string]cty.Value) err
 // Eval evaluates the parsed HCL documents using the input variables and populates v
 // using the result.
 func (s *State) Eval(parsed *hclparse.Parser, v any, input map[string]cty.Value) error {
-	ctx := s.newCtx()
-	reg := &blockDef{
-		fields:   make(map[string]struct{}),
-		children: make(map[string]*blockDef),
-	}
-	files := parsed.Files()
-	fileNames := make([]string, 0, len(files))
-	allBlocks := make([]*hclsyntax.Block, 0, len(files))
+	var (
+		ctx          = s.newCtx()
+		files        = parsed.Files()
+		fileNames    = make([]string, 0, len(files))
+		metaBlocks   = make(map[string][]*hclsyntax.Block, len(files))
+		staticBlocks = make([]*hclsyntax.Block, 0, len(files))
+		reg          = &blockDef{
+			fields:   make(map[string]struct{}),
+			children: make(map[string]*blockDef),
+		}
+	)
+	if ctx.Variables == nil {
+		ctx.Variables = make(map[string]cty.Value)
+	}
 	for name, file := range files {
 		fileNames = append(fileNames, name)
 		if err := s.setInputVals(ctx, file.Body, input); err != nil {
@@ -314,34 +320,57 @@ func (s *State) Eval(parsed *hclparse.Parser, v any, input map[string]cty.Value)
 		blocks := make(hclsyntax.Blocks, 0, len(body.Blocks))
 		for _, b := range body.Blocks {
 			switch {
-			// Variable blocks are not reachable by reference.
 			case b.Type == BlockVariable:
-				continue
-			// Semi-evaluate blocks with the for_each meta argument.
 			case b.Body != nil && b.Body.Attributes[forEachAttr] != nil:
-				nb, err := forEachBlocks(ctx, b)
-				if err != nil {
-					return err
-				}
-				blocks = append(blocks, nb...)
+				metaBlocks[name] = append(metaBlocks[name], b)
 			default:
 				blocks = append(blocks, b)
+				reg.child(extractDef(b, reg))
 			}
-			reg.child(extractDef(b, reg))
 		}
 		body.Blocks = blocks
-		allBlocks = append(allBlocks, blocks...)
+		staticBlocks = append(staticBlocks, blocks...)
 	}
-	vars, err := blockVars(allBlocks, "", reg)
+	vars, err := blockVars(staticBlocks, "", reg)
 	if err != nil {
 		return err
 	}
-	if ctx.Variables == nil {
-		ctx.Variables = make(map[string]cty.Value)
-	}
 	for k, v := range vars {
 		ctx.Variables[k] = v
 	}
+	// Semi-evaluate blocks with the for_each meta argument.
+	if len(metaBlocks) > 0 {
+		blocks := make([]*hclsyntax.Block, 0, len(metaBlocks))
+		for name, bs := range metaBlocks {
+			for _, b := range bs {
+				nb, err := forEachBlocks(ctx, b)
+				if err != nil {
+					return err
+				}
+				// Extract the definition of the top-level is enough.
+				reg.child(extractDef(b, reg))
+				blocks = append(blocks, nb...)
+				files[name].Body.(*hclsyntax.Body).Blocks = append(files[name].Body.(*hclsyntax.Body).Blocks, nb...)
+			}
+		}
+		if vars, err = blockVars(blocks, "", reg); err != nil {
+			return err
+		}
+		for k, v := range vars {
+			if v.IsNull() {
+				continue
+			}
+			if bs, ok := ctx.Variables[k]; !ok || bs.IsNull() {
+				ctx.Variables[k] = v
+			} else {
+				vs := bs.AsValueMap()
+				for k1, v1 := range v.AsValueMap() {
+					vs[k1] = v1
+				}
+				ctx.Variables[k] = cty.ObjectVal(vs)
+			}
+		}
+	}
 	spec := &Resource{}
 	sort.Slice(fileNames, func(i, j int) bool {
 		return fileNames[i] < fileNames[j]


Test patch:
diff --git a/schemahcl/schemahcl_test.go b/schemahcl/schemahcl_test.go
index 1928318d..e3ad75ac 100644
--- a/schemahcl/schemahcl_test.go
+++ b/schemahcl/schemahcl_test.go
@@ -461,6 +461,47 @@ env "dev" {
 		}),
 	})
 	require.EqualError(t, err, `variable "domains": a number is required`)
+
+	var (
+		// For-each resource depends on other resources.
+		doc1 struct {
+			Schema []*struct {
+				Name string `spec:",name"`
+			} `spec:"schema"`
+			Table []*struct {
+				Name   string `spec:"name,name"`
+				Schema *Ref   `spec:"schema"`
+			} `spec:"table"`
+		}
+		b1 = []byte(`
+schema "s1" {}
+schema "s2" {}
+
+table {
+  for_each = {
+    t1 = schema.s1
+	t2 = schema.s2
+  }
+  name = each.key
+  schema = each.value
+}
+`)
+	)
+	err = New().EvalBytes(b1, &doc1, nil)
+	require.NoError(t, err)
+	buf, err := Marshal.MarshalSpec(&doc1)
+	require.NoError(t, err)
+	require.Equal(t, `schema "s1" {
+}
+schema "s2" {
+}
+table "t1" {
+  schema = schema.s1
+}
+table "t2" {
+  schema = schema.s2
+}
+`, string(buf))
 }
 
 func TestDataLocalsRefs(t *testing.T) {