Repository: ariga/atlas
Pull request #1812 (2023-07-02 12:36:26)

schemahcl: skip evaluating non-referenced data-source


Patch:
diff --git a/cmd/atlas/internal/cmdapi/project.go b/cmd/atlas/internal/cmdapi/project.go
index 7c462b95..e9f2421f 100644
--- a/cmd/atlas/internal/cmdapi/project.go
+++ b/cmd/atlas/internal/cmdapi/project.go
@@ -451,11 +451,9 @@ func partialParse(path, env string) (*hclparse.Parser, error) {
 	if err != nil {
 		return nil, err
 	}
-	var used, datasrc []*hclsyntax.Block
+	var used []*hclsyntax.Block
 	for _, b := range fi.Body.(*hclsyntax.Body).Blocks {
 		switch b.Type {
-		case schemahcl.BlockData:
-			datasrc = append(datasrc, b)
 		case blockEnv:
 			switch n := len(b.Labels); {
 			// No env was selected.
@@ -474,14 +472,9 @@ func partialParse(path, env string) (*hclparse.Parser, error) {
 			used = append(used, b)
 		}
 	}
-	body := &hclsyntax.Body{
+	fi.Body = &hclsyntax.Body{
 		Blocks:     used,
 		Attributes: fi.Body.(*hclsyntax.Body).Attributes,
 	}
-	// If no attributes, locals, variables, or blocks reference data sources, we skip
-	// loading. Otherwise, body stays unchanged as data sources may depend on others.
-	if len(schemahcl.DynamicRefs(body).Data) == 0 {
-		fi.Body = body
-	}
 	return parser, nil
 }
diff --git a/schemahcl/context.go b/schemahcl/context.go
index 5fd2b464..c37a57b5 100644
--- a/schemahcl/context.go
+++ b/schemahcl/context.go
@@ -72,7 +72,7 @@ func (s *State) setInputVals(ctx *hcl.EvalContext, body hcl.Body, input map[stri
 	return nil
 }
 
-// evalReferences evaluates data blocks.
+// evalReferences evaluates local and data blocks.
 func (s *State) evalReferences(ctx *hcl.EvalContext, body *hclsyntax.Body) error {
 	type node struct {
 		addr  [3]string
@@ -141,7 +141,7 @@ func (s *State) evalReferences(ctx *hcl.EvalContext, body *hclsyntax.Body) error
 		}
 	}
 	var (
-		visit    func(n *node) error
+		visit    func(*node) error
 		visited  = make(map[*node]bool)
 		progress = make(map[*node]bool)
 	)
@@ -212,7 +212,25 @@ func (s *State) evalReferences(ctx *hcl.EvalContext, body *hclsyntax.Body) error
 			return err
 		}
 	}
+	dataref := dataRefs(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 {
+			exists := func() bool {
+				for _, r := range dataref {
+					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] {
+						return true
+					}
+				}
+				return false
+			}()
+			if !exists {
+				continue
+			}
+		}
 		if err := visit(n); err != nil {
 			return err
 		}
@@ -391,43 +409,6 @@ type Variables struct {
 	Data       map[string]map[string]bool
 }
 
-// DynamicRefs returns the references to dynamic variables.
-// i.e, variables, locals, and data sources.
-func DynamicRefs(b *hclsyntax.Body) *Variables {
-	var (
-		vars = &Variables{
-			Var:   make(map[string]bool),
-			Local: make(map[string]bool),
-			Data:  make(map[string]map[string]bool),
-		}
-		visit func(b *hclsyntax.Body, vars *Variables)
-	)
-	visit = func(b *hclsyntax.Body, vars *Variables) {
-		for _, a := range b.Attributes {
-			for _, v := range hclsyntax.Variables(a.Expr) {
-				switch {
-				case len(v) == 0:
-				case v.RootName() == RefVar && len(v) > 1:
-					vars.Var[v[1].(hcl.TraverseAttr).Name] = true
-				case v.RootName() == RefLocal && len(v) > 1:
-					vars.Local[v[1].(hcl.TraverseAttr).Name] = true
-				case v.RootName() == RefData && len(v) > 2:
-					n1, n2 := v[1].(hcl.TraverseAttr).Name, v[2].(hcl.TraverseAttr).Name
-					if _, ok := vars.Data[n1]; !ok {
-						vars.Data[n1] = make(map[string]bool)
-					}
-					vars.Data[n1][n2] = true
-				}
-			}
-		}
-		for _, b := range b.Blocks {
-			visit(b.Body, vars)
-		}
-	}
-	visit(b, vars)
-	return vars
-}
-
 // defRegistry returns a tree of blockDef structs representing the schema of the
 // blocks in the *hclsyntax.Body. The returned fields and children of each type
 // are an intersection of all existing blocks of the same type.
@@ -499,11 +480,28 @@ func extractDef(blk *hclsyntax.Block, parent *blockDef) *blockDef {
 }
 
 func bodyVars(b *hclsyntax.Body) (vars []hcl.Traversal) {
-	for _, attr := range b.Attributes {
-		vars = append(vars, hclsyntax.Variables(attr.Expr)...)
+	for _, a := range b.Attributes {
+		vars = append(vars, hclsyntax.Variables(a.Expr)...)
 	}
 	for _, b := range b.Blocks {
 		vars = append(vars, bodyVars(b.Body)...)
 	}
 	return
 }
+
+// dataRefs returns all data source referenced in the body.
+func dataRefs(b *hclsyntax.Body) (refs []hcl.Traversal) {
+	for _, a := range b.Attributes {
+		for _, v := range hclsyntax.Variables(a.Expr) {
+			if v.RootName() == RefData {
+				refs = append(refs, v)
+			}
+		}
+	}
+	for _, b := range b.Blocks {
+		if b.Type != BlockData {
+			refs = append(refs, dataRefs(b.Body)...)
+		}
+	}
+	return
+}


Test patch:
diff --git a/cmd/atlas/internal/cmdext/cmdext_test.go b/cmd/atlas/internal/cmdext/cmdext_test.go
index a8f08fe5..fa381c0e 100644
--- a/cmd/atlas/internal/cmdext/cmdext_test.go
+++ b/cmd/atlas/internal/cmdext/cmdext_test.go
@@ -440,6 +440,8 @@ func TestRemoteDir(t *testing.T) {
 data "remote_dir" "hello" {
   name  = "atlas"
 }
+
+dir = data.remote_dir.hello.url
 `), &v, map[string]cty.Value{"cloud_url": cty.StringVal(srv.URL)})
 	require.EqualError(t, err, "data.remote_dir.hello: missing atlas cloud config")
 
diff --git a/schemahcl/context_test.go b/schemahcl/context_test.go
index 4e3c5eb0..2b3fa45b 100644
--- a/schemahcl/context_test.go
+++ b/schemahcl/context_test.go
@@ -8,8 +8,6 @@ import (
 	"reflect"
 	"testing"
 
-	"github.com/hashicorp/hcl/v2/hclparse"
-	"github.com/hashicorp/hcl/v2/hclsyntax"
 	"github.com/stretchr/testify/require"
 	"github.com/zclconf/go-cty/cty"
 )
@@ -614,77 +612,3 @@ variable "name" {
 	err = New().EvalBytes([]byte(h), &struct{}{}, nil)
 	require.EqualError(t, err, `invalid type "boring" for variable "name". Valid types are: string, number, bool, list, map, or set`)
 }
-
-func TestDynamicRefs(t *testing.T) {
-	var (
-		h = `
-variable "name" {
-  type = string
-}
-
-locals {
-  name = "a8m"
-}
-
-data "remote_dir" "example1" {
-  name = "ent"
-}
-
-data "remote_dir" "example2" {
-  name = "atlas"
-}
-
-data "template_dir" "example1" {
-  name = "atlas"
-}
-
-env "var_only" {
-  v = var.name
-}
-
-env "local_only" {
-  v = local.name
-}
-
-env "data_only" {
-  v = data.remote_dir.example1.url
-}
-
-env "var_local" {
-  v = "${var.name}-${local.name}"
-}
-
-env "var_data" {
-  v = "${var.name}-${data.remote_dir.example1.url}"
-}
-
-env "all" {
-  v = "${var.name}-${local.name}-${data.remote_dir.example1.url}-${data.remote_dir.example2.url}-${data.template_dir.example1.url}"
-}
-`
-		zero1, zero2 = make(map[string]bool), make(map[string]map[string]bool)
-	)
-	for name, vars := range map[string]Variables{
-		"var_only":   {Var: map[string]bool{"name": true}, Local: zero1, Data: zero2},
-		"local_only": {Var: zero1, Local: map[string]bool{"name": true}, Data: zero2},
-		"data_only":  {Var: zero1, Local: zero1, Data: map[string]map[string]bool{"remote_dir": {"example1": true}}},
-		"var_local":  {Var: map[string]bool{"name": true}, Local: map[string]bool{"name": true}, Data: zero2},
-		"var_data":   {Var: map[string]bool{"name": true}, Local: zero1, Data: map[string]map[string]bool{"remote_dir": {"example1": true}}},
-		"all":        {Var: map[string]bool{"name": true}, Local: map[string]bool{"name": true}, Data: map[string]map[string]bool{"remote_dir": {"example1": true, "example2": true}, "template_dir": {"example1": true}}},
-	} {
-		parser := hclparse.NewParser()
-		fi, diags := parser.ParseHCL([]byte(h), "test")
-		require.False(t, diags.HasErrors())
-		env := func() *hclsyntax.Body {
-			for _, b := range fi.Body.(*hclsyntax.Body).Blocks {
-				if b.Type == "env" && b.Labels[0] == name {
-					return b.Body
-				}
-			}
-			require.Failf(t, "env %q block not found", name)
-			return nil
-		}()
-		got := DynamicRefs(env)
-		require.Equal(t, vars, *got)
-	}
-}
diff --git a/schemahcl/hcl_test.go b/schemahcl/hcl_test.go
index 1e703c2e..47431685 100644
--- a/schemahcl/hcl_test.go
+++ b/schemahcl/hcl_test.go
@@ -639,3 +639,62 @@ locals {
 `)
 	require.Error(t, New(opts...).EvalBytes(b, &doc, nil), `cyclic reference to "data.text.a"`)
 }
+
+func TestSkippedDataSrc(t *testing.T) {
+	var (
+		opts = []Option{
+			WithDataSource("dynamic", func(ctx *hcl.EvalContext, b *hclsyntax.Block) (cty.Value, error) {
+				attrs, diags := b.Body.JustAttributes()
+				if diags.HasErrors() {
+					return cty.NilVal, diags
+				}
+				s, diags := attrs["skip"].Expr.Value(ctx)
+				if diags.HasErrors() {
+					return cty.NilVal, diags
+				}
+				if s.True() {
+					return cty.NilVal, fmt.Errorf("data source should be skipped, but was called with %q", b.Labels)
+				}
+				v, diags := attrs["v"].Expr.Value(ctx)
+				if diags.HasErrors() {
+					return cty.NilVal, diags
+				}
+				return cty.ObjectVal(map[string]cty.Value{
+					"v": v,
+				}), nil
+			}),
+		}
+		v struct {
+			V2 string `spec:"v2"`
+			V3 string `spec:"v3"`
+		}
+		b = []byte(`
+data "dynamic" "skipped1" {
+  v = "value is irrelevant"
+  // This attribute has no meaning, besides indicating
+  // to the test that the data source should be skipped.
+  skip = true
+}
+data "dynamic" "skipped2" {
+  v = data.dynamic.skipped1.v
+  skip = true
+}
+data "dynamic" "evaluated1" {
+  v = "v1"
+  skip = false
+}
+data "dynamic" "evaluated2" {
+  v = "v2"
+  skip = false
+}
+data "dynamic" "evaluated3" {
+  v = data.dynamic.evaluated1.v
+  skip = false
+}
+
+v2 = data.dynamic.evaluated2.v
+v3 = data.dynamic.evaluated3.v
+`)
+	)
+	require.NoError(t, New(opts...).EvalBytes(b, &v, nil))
+}