Repository: ariga/atlas
Pull request #2309 (2023-11-22 22:35:36)

schemahcl: allow registering custom type-block to first eval phase


Patch:
diff --git a/schemahcl/context.go b/schemahcl/context.go
index be60bc0a..085e9bb4 100644
--- a/schemahcl/context.go
+++ b/schemahcl/context.go
@@ -81,6 +81,7 @@ func (s *State) evalReferences(ctx *hcl.EvalContext, body *hclsyntax.Body) error
 	}
 	var (
 		initblk []*node
+		typeblk = make(map[string]bool)
 		nodes   = make(map[[3]string]*node)
 		blocks  = make(hclsyntax.Blocks, 0, len(body.Blocks))
 	)
@@ -136,6 +137,24 @@ func (s *State) evalReferences(ctx *hcl.EvalContext, body *hclsyntax.Body) error
 			}
 			nodes[addr] = n
 			initblk = append(initblk, n)
+		case s.config.typedblk[b.Type] != nil:
+			typeblk[b.Type] = true
+			if len(b.Labels) < 2 {
+				return fmt.Errorf("%s block must have exactly 2 labels", b.Type)
+			}
+			k, ok := s.config.typedblk[b.Type]
+			if !ok || k[b.Labels[0]] == nil {
+				return fmt.Errorf("missing %s block handler for %q", b.Type, b.Labels[0])
+			}
+			h := k[b.Labels[0]]
+			// Typed block references are combined from
+			// "<type>", "<label>" and "name" labels.
+			addr := [3]string{b.Type, b.Labels[0], b.Labels[1]}
+			nodes[addr] = &node{
+				addr:  addr,
+				value: func() (cty.Value, error) { return h(ctx, b) },
+				edges: func() []hcl.Traversal { return bodyVars(b.Body) },
+			}
 		default:
 			blocks = append(blocks, b)
 		}
@@ -162,8 +181,8 @@ func (s *State) evalReferences(ctx *hcl.EvalContext, body *hclsyntax.Body) error
 			switch root := e.RootName(); {
 			case root == RefLocal && len(e) > 1:
 				addr = [3]string{RefLocal, e[1].(hcl.TraverseAttr).Name, ""}
-			case root == RefData && len(e) > 2:
-				addr = [3]string{RefData, e[1].(hcl.TraverseAttr).Name, e[2].(hcl.TraverseAttr).Name}
+			case (root == RefData || typeblk[e.RootName()]) && len(e) > 2:
+				addr = [3]string{e.RootName(), e[1].(hcl.TraverseAttr).Name, e[2].(hcl.TraverseAttr).Name}
 			case s.config.initblk[root] != nil && len(e) == 1:
 				addr = [3]string{root, "", ""}
 			}
@@ -180,20 +199,20 @@ func (s *State) evalReferences(ctx *hcl.EvalContext, body *hclsyntax.Body) error
 		if err != nil {
 			return err
 		}
-		switch n.addr[0] {
-		case RefData:
-			data := make(map[string]cty.Value)
-			if vv, ok := ctx.Variables[RefData]; ok {
-				data = vv.AsValueMap()
+		switch {
+		case n.addr[0] == RefData, typeblk[n.addr[0]]:
+			top := make(map[string]cty.Value)
+			if vv, ok := ctx.Variables[n.addr[0]]; ok {
+				top = vv.AsValueMap()
 			}
-			src := make(map[string]cty.Value)
-			if vv, ok := data[n.addr[1]]; ok {
-				src = vv.AsValueMap()
+			typ := make(map[string]cty.Value)
+			if vv, ok := top[n.addr[1]]; ok {
+				typ = vv.AsValueMap()
 			}
-			src[n.addr[2]] = v
-			data[n.addr[1]] = cty.ObjectVal(src)
-			ctx.Variables[RefData] = cty.ObjectVal(data)
-		case RefLocal:
+			typ[n.addr[2]] = v
+			top[n.addr[1]] = cty.ObjectVal(typ)
+			ctx.Variables[n.addr[0]] = cty.ObjectVal(top)
+		case n.addr[0] == RefLocal:
 			locals := make(map[string]cty.Value)
 			if vv, ok := ctx.Variables[RefLocal]; ok {
 				locals = vv.AsValueMap()
@@ -212,13 +231,21 @@ func (s *State) evalReferences(ctx *hcl.EvalContext, body *hclsyntax.Body) error
 			return err
 		}
 	}
-	dataref := dataRefs(body)
+	// New body without locals, data-sources, typed-blocks attached.
+	body.Blocks = blocks
+	typeref := map[string][]hcl.Traversal{
+		RefData: typeRefs(RefData, body),
+	}
+	for n := range typeblk {
+		typeref[n] = typeRefs(n, body)
+	}
 	for _, n := range nodes {
-		// Evaluate data sources only if they were referenced by other top-level
-		// blocks/attributes or if they reference other evaluated data sources.
-		if n.addr[0] == BlockData {
+		// Evaluate type-blocks and data sources only if they were referenced by
+		// other top-level blocks/attributes or if they reference other evaluated
+		// data sources.
+		if refs, ok := typeref[n.addr[0]]; ok {
 			exists := func() bool {
-				for _, r := range dataref {
+				for _, r := range refs {
 					t, ok1 := r[1].(hcl.TraverseAttr)
 					l, ok2 := r[2].(hcl.TraverseAttr)
 					if ok1 && ok2 && t.Name == n.addr[1] && l.Name == n.addr[2] {
@@ -235,7 +262,6 @@ func (s *State) evalReferences(ctx *hcl.EvalContext, body *hclsyntax.Body) error
 			return err
 		}
 	}
-	body.Blocks = blocks
 	return nil
 }
 
@@ -494,18 +520,18 @@ func bodyVars(b *hclsyntax.Body) (vars []hcl.Traversal) {
 	return
 }
 
-// dataRefs returns all data source referenced in the body.
-func dataRefs(b *hclsyntax.Body) (refs []hcl.Traversal) {
+// typeRefs returns all type referenced in the body.
+func typeRefs(name string, b *hclsyntax.Body) (refs []hcl.Traversal) {
 	for _, a := range b.Attributes {
 		for _, v := range hclsyntax.Variables(a.Expr) {
-			if v.RootName() == RefData {
+			if v.RootName() == name {
 				refs = append(refs, v)
 			}
 		}
 	}
 	for _, b := range b.Blocks {
-		if b.Type != BlockData {
-			refs = append(refs, dataRefs(b.Body)...)
+		if b.Type != name {
+			refs = append(refs, typeRefs(name, b.Body)...)
 		}
 	}
 	return
diff --git a/schemahcl/schemahcl.go b/schemahcl/schemahcl.go
index 1b16979d..c51ce067 100644
--- a/schemahcl/schemahcl.go
+++ b/schemahcl/schemahcl.go
@@ -29,8 +29,9 @@ type (
 		funcs            map[string]function.Function
 		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)
 		validator        func() SchemaValidator
+		datasrc, initblk map[string]func(*hcl.EvalContext, *hclsyntax.Block) (cty.Value, error)
+		typedblk         map[string]map[string]func(*hcl.EvalContext, *hclsyntax.Block) (cty.Value, error)
 	}
 	// Option configures a Config.
 	Option func(*Config)
@@ -137,6 +138,37 @@ func WithDataSource(name string, h func(*hcl.EvalContext, *hclsyntax.Block) (cty
 	}
 }
 
+// WithTypeLabelBlock registers a type-block and its label along with the corresponding handler.
+// e.g., the example below registers a typed block named "driver" with the label "remote" that
+// returns the string defined in the token attribute.
+//
+//	WithTypeLabelBlock("driver", "remote", 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["token"].Expr.Value(ctx)
+//		if diags.HasErrors() {
+//			return cty.NilVal, diags
+//		}
+//		return cty.ObjectVal(map[string]cty.Value{"url": v}), nil
+//	})
+//
+//	driver "remote" "hello" {
+//	  token = "hello world"
+//	}
+func WithTypeLabelBlock(name, label string, h func(*hcl.EvalContext, *hclsyntax.Block) (cty.Value, error)) Option {
+	return func(c *Config) {
+		if c.typedblk == nil {
+			c.typedblk = make(map[string]map[string]func(*hcl.EvalContext, *hclsyntax.Block) (cty.Value, error))
+		}
+		if c.typedblk[name] == nil {
+			c.typedblk[name] = make(map[string]func(*hcl.EvalContext, *hclsyntax.Block) (cty.Value, error))
+		}
+		c.typedblk[name][label] = h
+	}
+}
+
 // WithInitBlock registers a block that evaluates (first) to a cty.Value,
 // has no labels, and can be defined only once. For example:
 //
@@ -542,8 +574,8 @@ func (s *State) typeError(diag hcl.Diagnostics, scope []string) error {
 			if d.Summary != "Unknown variable" {
 				continue
 			}
-			switch root := e.Traversal.RootName(); root {
-			case RefData:
+			switch root := e.Traversal.RootName(); {
+			case root == RefData, s.config.typedblk[root] != nil:
 				var b strings.Builder
 				b.WriteString(root)
 				for _, t := range e.Traversal[1:] {
@@ -553,8 +585,11 @@ func (s *State) typeError(diag hcl.Diagnostics, scope []string) error {
 					}
 				}
 				d.Summary = "Unknown data source"
+				if s.config.typedblk[root] != nil {
+					d.Summary = "Unknown block type"
+				}
 				d.Detail = fmt.Sprintf("%s does not exist", b.String())
-			case RefLocal:
+			case root == RefLocal:
 				d.Summary = "Unknown local"
 			default:
 				if t, ok := s.findTypeSpec(root); ok && len(t.Attributes) > 0 {


Test patch:
diff --git a/schemahcl/schemahcl_test.go b/schemahcl/schemahcl_test.go
index 29adaab9..0c43ac8e 100644
--- a/schemahcl/schemahcl_test.go
+++ b/schemahcl/schemahcl_test.go
@@ -702,6 +702,158 @@ v3 = data.dynamic.evaluated3.v
 `)
 	)
 	require.NoError(t, New(opts...).EvalBytes(b, &v, nil))
+
+	b = []byte(`
+locals {
+  a = data.dynamic.not_skipped1.v
+  b = local.c
+  c = "3"
+  d = data.dynamic.not_skipped3.v
+}
+
+data "dynamic" "not_skipped1" {
+  v = "v2"
+  skip = false
+}
+
+data "dynamic" "not_skipped2" {
+  v = "v${local.b}"
+  skip = false
+}
+
+data "dynamic" "not_skipped3" {
+  v = "v4"
+  skip = false
+}
+
+data "dynamic" "not_skipped4" {
+  v = "v5"
+  skip = false
+}
+
+data "dynamic" "skipped1" {
+  v = local.a
+  skip = true
+}
+
+data "dynamic" "skipped2" {
+  v = "v"
+  skip = true
+}
+
+data "dynamic" "skipped3" {
+  v = data.dynamic.skipped2
+  skip = true
+}
+
+top {
+  a = data.dynamic.not_skipped2.v // "v3".
+  block {
+    block {
+      a1 = local.d // "v4".
+      a2 = data.dynamic.not_skipped4.v // "v5".
+    }
+  }
+}
+
+v2 = local.a
+`)
+	var v1 struct {
+		Top struct {
+			A     string `spec:"a"`
+			Block struct {
+				Block struct {
+					A1 string `spec:"a1"`
+					A2 string `spec:"a2"`
+				} `spec:"block"`
+			} `spec:"block"`
+		} `spec:"top"`
+		V2 string `spec:"v2"`
+	}
+	require.NoError(t, New(opts...).EvalBytes(b, &v1, nil))
+	require.Equal(t, "v2", v1.V2)
+	require.Equal(t, "v3", v1.Top.A)
+	require.Equal(t, "v4", v1.Top.Block.Block.A1)
+	require.Equal(t, "v5", v1.Top.Block.Block.A2)
+}
+
+func TestTypeLabelBlock(t *testing.T) {
+	var (
+		callD, callT int
+		opts         = []Option{
+			WithTypeLabelBlock("driver", "remote", func(ctx *hcl.EvalContext, b *hclsyntax.Block) (cty.Value, error) {
+				attrs, diags := b.Body.JustAttributes()
+				if diags.HasErrors() {
+					return cty.NilVal, diags
+				}
+				v, diags := attrs["name"].Expr.Value(ctx)
+				if diags.HasErrors() {
+					return cty.NilVal, diags
+				}
+				callT++
+				return cty.ObjectVal(map[string]cty.Value{"url": cty.StringVal("driver://" + v.AsString())}), nil
+			}),
+			WithTypeLabelBlock("driver", "not_called", func(ctx *hcl.EvalContext, b *hclsyntax.Block) (cty.Value, error) {
+				t.Fatal("should not be called")
+				return cty.NilVal, nil
+			}),
+			WithDataSource("text", func(ctx *hcl.EvalContext, b *hclsyntax.Block) (cty.Value, error) {
+				attrs, diags := b.Body.JustAttributes()
+				if diags.HasErrors() {
+					return cty.NilVal, diags
+				}
+				v, diags := attrs["value"].Expr.Value(ctx)
+				if diags.HasErrors() {
+					return cty.NilVal, diags
+				}
+				callD++
+				return cty.ObjectVal(map[string]cty.Value{"output": v}), nil
+			}),
+		}
+		doc struct {
+			Values []string `spec:"vs"`
+		}
+		b = []byte(`
+locals {
+  a = "a8m"
+}
+
+data "text" "a" {
+  value = local.a
+}
+
+driver "remote" "myapp" {
+  name = data.text.a.output
+}
+
+vs = [
+  driver.remote.myapp.url,
+  data.text.a.output
+]
+`)
+	)
+	require.NoError(t, New(opts...).EvalBytes(b, &doc, nil))
+	require.Equal(t, []string{"driver://a8m", "a8m"}, doc.Values)
+	require.EqualValues(t, 1, callT)
+	require.EqualValues(t, 2, callD, "it is up to the data source to implement caching")
+
+	b = []byte(`
+locals {
+  a = "a8m"
+}
+
+data "text" "a" {
+  value = local.a
+}
+
+driver "remote" "myapp" {
+  name = data.text.a.output
+}
+`)
+	require.NoError(t, New(opts...).EvalBytes(b, &doc, nil))
+	require.Equal(t, []string{"driver://a8m", "a8m"}, doc.Values)
+	require.Equal(t, 1, callT)
+	require.Equal(t, 2, callD)
 }
 
 type countValidator struct{ nb, na int }