Repository: ariga/atlas
Pull request #1507 (2023-03-04 21:24:40)

schemahcl: add support for init blocks


Patch:
diff --git a/schemahcl/context.go b/schemahcl/context.go
index 06378a8e..1ea38708 100644
--- a/schemahcl/context.go
+++ b/schemahcl/context.go
@@ -80,12 +80,13 @@ func (s *State) evalReferences(ctx *hcl.EvalContext, body *hclsyntax.Body) error
 		value func() (cty.Value, error)
 	}
 	var (
-		nodes  = make(map[[3]string]*node)
-		blocks = make(hclsyntax.Blocks, 0, len(body.Blocks))
+		initblk []*node
+		nodes   = make(map[[3]string]*node)
+		blocks  = make(hclsyntax.Blocks, 0, len(body.Blocks))
 	)
 	for _, b := range body.Blocks {
-		switch b := b; b.Type {
-		case dataBlock:
+		switch b := b; {
+		case b.Type == dataBlock:
 			if len(b.Labels) < 2 {
 				return fmt.Errorf("data block %q must have exactly 2 labels", b.Type)
 			}
@@ -101,7 +102,7 @@ func (s *State) evalReferences(ctx *hcl.EvalContext, body *hclsyntax.Body) error
 				value: func() (cty.Value, error) { return h(ctx, b) },
 				edges: func() []hcl.Traversal { return bodyVars(b.Body) },
 			}
-		case localsBlock:
+		case b.Type == localsBlock:
 			for k, v := range b.Body.Attributes {
 				k, v := k, v
 				// Local references are combined from
@@ -119,6 +120,22 @@ func (s *State) evalReferences(ctx *hcl.EvalContext, body *hclsyntax.Body) error
 					},
 				}
 			}
+		case s.config.initblk[b.Type] != nil:
+			if len(b.Labels) != 0 {
+				return fmt.Errorf("init block %q cannot have labels", b.Type)
+			}
+			addr := [3]string{b.Type, "", ""}
+			if nodes[addr] != nil {
+				return fmt.Errorf("duplicate init block %q", b.Type)
+			}
+			h := s.config.initblk[b.Type]
+			n := &node{
+				addr:  addr,
+				value: func() (cty.Value, error) { return h(ctx, b) },
+				edges: func() []hcl.Traversal { return bodyVars(b.Body) },
+			}
+			nodes[addr] = n
+			initblk = append(initblk, n)
 		default:
 			blocks = append(blocks, b)
 		}
@@ -134,8 +151,8 @@ func (s *State) evalReferences(ctx *hcl.EvalContext, body *hclsyntax.Body) error
 		}
 		if progress[n] {
 			addr := n.addr[:]
-			if addr[2] == "" {
-				addr = addr[:2]
+			for len(addr) > 0 && addr[len(addr)-1] == "" {
+				addr = addr[:len(addr)-1]
 			}
 			return fmt.Errorf("cyclic reference to %q", strings.Join(addr, "."))
 		}
@@ -147,6 +164,8 @@ func (s *State) evalReferences(ctx *hcl.EvalContext, body *hclsyntax.Body) error
 				addr = [3]string{localRef, e[1].(hcl.TraverseAttr).Name, ""}
 			case root == dataBlock && len(e) > 2:
 				addr = [3]string{dataBlock, e[1].(hcl.TraverseAttr).Name, e[2].(hcl.TraverseAttr).Name}
+			case s.config.initblk[root] != nil && len(e) == 1:
+				addr = [3]string{root, "", ""}
 			}
 			// Unrecognized reference.
 			if nodes[addr] == nil {
@@ -181,9 +200,18 @@ func (s *State) evalReferences(ctx *hcl.EvalContext, body *hclsyntax.Body) error
 			}
 			locals[n.addr[1]] = v
 			ctx.Variables[localRef] = cty.ObjectVal(locals)
+		default:
+			ctx.Variables[n.addr[0]] = v
 		}
 		return nil
 	}
+	// Evaluate init-blocks first,
+	// to give them higher precedence.
+	for _, n := range initblk {
+		if err := visit(n); err != nil {
+			return err
+		}
+	}
 	for _, n := range nodes {
 		if err := visit(n); err != nil {
 			return err
diff --git a/schemahcl/opts.go b/schemahcl/opts.go
index 2dcd652e..3e1fd4bf 100644
--- a/schemahcl/opts.go
+++ b/schemahcl/opts.go
@@ -18,11 +18,11 @@ import (
 type (
 	// Config configures an unmarshaling.
 	Config struct {
-		types     []*TypeSpec
-		newCtx    func() *hcl.EvalContext
-		pathVars  map[string]map[string]cty.Value
-		pathFuncs map[string]map[string]function.Function
-		datasrc   map[string]func(*hcl.EvalContext, *hclsyntax.Block) (cty.Value, error)
+		types            []*TypeSpec
+		newCtx           func() *hcl.EvalContext
+		pathVars         map[string]map[string]cty.Value
+		pathFuncs        map[string]map[string]function.Function
+		datasrc, initblk map[string]func(*hcl.EvalContext, *hclsyntax.Block) (cty.Value, error)
 	}
 	// Option configures a Config.
 	Option func(*Config)
@@ -96,7 +96,7 @@ func WithVariables(vars map[string]cty.Value) Option {
 //			return cty.NilVal, diags
 //		}
 //		return cty.ObjectVal(map[string]cty.Value{"output": v}), nil
-//	}),
+//	})
 //
 //	data "text" "hello" {
 //	  value = "hello world"
@@ -110,6 +110,29 @@ func WithDataSource(name string, h func(*hcl.EvalContext, *hclsyntax.Block) (cty
 	}
 }
 
+// WithInitBlock registers a block that evaluates (first) to a cty.Value,
+// has no labels, and can be defined only once. For example:
+//
+//	WithInitBlock("atlas", func(ctx *hcl.EvalContext, b *hclsyntax.Block) (cty.Value, hcl.Diagnostics) {
+//		attrs, diags := b.Body.JustAttributes()
+//		if diags.HasErrors() {
+//			return cty.NilVal, diags
+//		}
+//		v, diags := attrs["modules"].Expr.Value(ctx)
+//		if diags.HasErrors() {
+//			return cty.NilVal, diags
+//		}
+//		return cty.ObjectVal(map[string]cty.Value{"modules": v}), nil
+//	})
+func WithInitBlock(name string, h func(*hcl.EvalContext, *hclsyntax.Block) (cty.Value, error)) Option {
+	return func(c *Config) {
+		if c.initblk == nil {
+			c.initblk = make(map[string]func(*hcl.EvalContext, *hclsyntax.Block) (cty.Value, error))
+		}
+		c.initblk[name] = h
+	}
+}
+
 // WithTypes configures the given types as identifiers in the unmarshal
 // context. The path controls where the usage of this type is allowed.
 func WithTypes(path string, typeSpecs []*TypeSpec) Option {


Test patch:
diff --git a/schemahcl/hcl_test.go b/schemahcl/hcl_test.go
index d98ab405..f1b0564f 100644
--- a/schemahcl/hcl_test.go
+++ b/schemahcl/hcl_test.go
@@ -5,9 +5,11 @@
 package schemahcl
 
 import (
+	"errors"
 	"fmt"
 	"log"
 	"os"
+	"path"
 	"path/filepath"
 	"testing"
 
@@ -478,6 +480,51 @@ func TestDataLocalsRefs(t *testing.T) {
 				}
 				return cty.ObjectVal(map[string]cty.Value{"output": v}), nil
 			}),
+			WithInitBlock("atlas", func(ctx *hcl.EvalContext, b *hclsyntax.Block) (cty.Value, error) {
+				org, diags := b.Body.Attributes["org"].Expr.Value(ctx)
+				if diags.HasErrors() {
+					return cty.NilVal, diags
+				}
+				if len(b.Body.Blocks) != 1 || b.Body.Blocks[0].Type != "auth" {
+					return cty.NilVal, errors.New("expected auth block")
+				}
+				attrs, diags := b.Body.Blocks[0].Body.JustAttributes()
+				if diags.HasErrors() {
+					return cty.NilVal, diags
+				}
+				host, diags := attrs["host"].Expr.Value(ctx)
+				if diags.HasErrors() {
+					return cty.NilVal, diags
+				}
+				return cty.ObjectVal(map[string]cty.Value{
+					"org": org,
+					"auth": cty.ObjectVal(map[string]cty.Value{
+						"host": host,
+					}),
+				}), nil
+			}),
+			WithDataSource("remote_dir", func(ctx *hcl.EvalContext, b *hclsyntax.Block) (cty.Value, error) {
+				attrs, diags := b.Body.JustAttributes()
+				if diags.HasErrors() {
+					return cty.NilVal, diags
+				}
+				host, diags := attrs["host"].Expr.Value(ctx)
+				if diags.HasErrors() {
+					return cty.NilVal, diags
+				}
+				org, diags := (&hclsyntax.ScopeTraversalExpr{
+					Traversal: hcl.Traversal{
+						hcl.TraverseRoot{Name: "atlas", SrcRange: b.Range()},
+						hcl.TraverseAttr{Name: "org", SrcRange: b.Range()},
+					},
+				}).Value(ctx)
+				if diags.HasErrors() {
+					return cty.NilVal, diags
+				}
+				return cty.ObjectVal(map[string]cty.Value{
+					"url": cty.StringVal("atlas://" + path.Join(host.AsString(), org.AsString(), b.Labels[1])),
+				}), nil
+			}),
 		}
 		doc struct {
 			Values []string `spec:"vs"`
@@ -495,6 +542,7 @@ locals {
   // locals can reference data sources.
   c = "local-c-ref-data-a: ${data.text.a.output}"
   d = "local-d"
+  host = "atlasgo.io"
 }
 
 data "sql" "tenants" {
@@ -519,6 +567,17 @@ data "text" "b" {
   value = "data-text-b-ref-local-d: ${local.d}"
 }
 
+atlas {
+  org = "ent"
+  auth {
+    host = local.host
+  }
+}
+
+data "remote_dir" "migrations" {
+  host = atlas.auth.host
+}
+
 vs = [
   local.a,
   local.b,
@@ -526,6 +585,7 @@ vs = [
   data.sql.tenants.output,
   data.text.a.output,
   data.text.b.output,
+  data.remote_dir.migrations.url,
 ]
 `)
 	)
@@ -537,6 +597,7 @@ vs = [
 		"data-sql-tenants",
 		"data-text-a-ref-data-sql-tenants: data-sql-tenants",
 		"data-text-b-ref-local-d: local-d",
+		"atlas://atlasgo.io/ent/migrations",
 	}, doc.Values)
 
 	b = []byte(`locals { a = local.a }`)