Repository: ariga/atlas
Pull request #1670 (2023-05-18 10:41:43)

cmd/atlas: only evaluate necessary blocks in atlas.hcl


Patch:
diff --git a/cmd/atlas/internal/cmdapi/project.go b/cmd/atlas/internal/cmdapi/project.go
index 23066191..adf984d7 100644
--- a/cmd/atlas/internal/cmdapi/project.go
+++ b/cmd/atlas/internal/cmdapi/project.go
@@ -15,11 +15,12 @@ import (
 	"ariga.io/atlas/schemahcl"
 	"ariga.io/atlas/sql/schema"
 
+	"github.com/hashicorp/hcl/v2"
+	"github.com/hashicorp/hcl/v2/hclparse"
+	"github.com/hashicorp/hcl/v2/hclsyntax"
 	"github.com/zclconf/go-cty/cty"
 )
 
-const defaultConfigPath = "file://atlas.hcl"
-
 type loadConfig struct {
 	inputValues map[string]cty.Value
 }
@@ -385,7 +386,17 @@ func EnvByName(name string, opts ...LoadOption) (*Project, []*Env, error) {
 	return project, selected, nil
 }
 
+const (
+	blockEnv          = "env"
+	refAtlas          = "atlas"
+	defaultConfigPath = "file://atlas.hcl"
+)
+
 func parseConfig(path, env string, values map[string]cty.Value) (*Project, error) {
+	pr, err := partialParse(path, env)
+	if err != nil {
+		return nil, err
+	}
 	cfg := &cmdext.AtlasConfig{}
 	opts := append(
 		cmdext.DataSources,
@@ -397,13 +408,13 @@ func parseConfig(path, env string, values map[string]cty.Value) (*Project, error
 			formatGolangMigrate,
 		),
 		schemahcl.WithVariables(map[string]cty.Value{
-			"atlas": cty.ObjectVal(map[string]cty.Value{
-				"env": cty.StringVal(env),
+			refAtlas: cty.ObjectVal(map[string]cty.Value{
+				blockEnv: cty.StringVal(env),
 			}),
 		}),
 	)
 	p := &Project{Lint: &Lint{}, Diff: &Diff{}, cfg: cfg}
-	if err := schemahcl.New(opts...).EvalFiles([]string{path}, p, values); err != nil {
+	if err := schemahcl.New(opts...).Eval(pr, p, values); err != nil {
 		return nil, err
 	}
 	for _, e := range p.Envs {
@@ -413,5 +424,57 @@ func parseConfig(path, env string, values map[string]cty.Value) (*Project, error
 }
 
 func init() {
-	schemahcl.Register("env", &Env{})
+	schemahcl.Register(blockEnv, &Env{})
+}
+
+func partialParse(path, env string) (*hclparse.Parser, error) {
+	parser := hclparse.NewParser()
+	fi, err := parser.ParseHCLFile(path)
+	if err != nil {
+		return nil, err
+	}
+	var used, locals, datasrc []*hclsyntax.Block
+	for _, b := range fi.Body.(*hclsyntax.Body).Blocks {
+		switch b.Type {
+		case schemahcl.BlockLocals:
+			locals = append(locals, b)
+		case schemahcl.BlockData:
+			datasrc = append(datasrc, b)
+		case blockEnv:
+			switch n := len(b.Labels); {
+			case n == 1 && b.Labels[0] == env:
+				used = append(used, b)
+			case n == 0 && b.Body != nil && b.Body.Attributes[schemahcl.AttrName] != nil:
+				t, ok := b.Body.Attributes[schemahcl.AttrName].Expr.(*hclsyntax.ScopeTraversalExpr)
+				if ok && len(t.Traversal) == 2 && t.Traversal.RootName() == refAtlas && t.Traversal[1].(hcl.TraverseAttr).Name == blockEnv {
+					used = append(used, b)
+				}
+			}
+		default:
+			used = append(used, b)
+		}
+	}
+	var (
+		body = &hclsyntax.Body{
+			Blocks:     used,
+			Attributes: fi.Body.(*hclsyntax.Body).Attributes,
+		}
+		vars = schemahcl.DynamicRefs(body)
+	)
+	for _, l := range locals {
+		for _, a := range l.Body.Attributes {
+			if !vars.Local[a.Name] {
+				// Remove local variables that are not used.
+				delete(l.Body.Attributes, a.Name)
+			}
+		}
+	}
+	body.Blocks = append(body.Blocks, locals...)
+	for _, d := range datasrc {
+		if len(d.Labels) == 2 && vars.Data[d.Labels[0]] != nil && vars.Data[d.Labels[0]][d.Labels[1]] {
+			body.Blocks = append(body.Blocks, d)
+		}
+	}
+	fi.Body = body
+	return parser, nil
 }
diff --git a/schemahcl/context.go b/schemahcl/context.go
index 1ea38708..971f04fb 100644
--- a/schemahcl/context.go
+++ b/schemahcl/context.go
@@ -86,7 +86,7 @@ func (s *State) evalReferences(ctx *hcl.EvalContext, body *hclsyntax.Body) error
 	)
 	for _, b := range body.Blocks {
 		switch b := b; {
-		case b.Type == dataBlock:
+		case b.Type == BlockData:
 			if len(b.Labels) < 2 {
 				return fmt.Errorf("data block %q must have exactly 2 labels", b.Type)
 			}
@@ -96,18 +96,18 @@ func (s *State) evalReferences(ctx *hcl.EvalContext, body *hclsyntax.Body) error
 			}
 			// Data references are combined from
 			// "data", "source" and "name" labels.
-			addr := [3]string{dataBlock, b.Labels[0], b.Labels[1]}
+			addr := [3]string{RefData, 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) },
 			}
-		case b.Type == localsBlock:
+		case b.Type == BlockLocals:
 			for k, v := range b.Body.Attributes {
 				k, v := k, v
 				// Local references are combined from
 				// "local" and "name" labels.
-				addr := [3]string{localRef, k, ""}
+				addr := [3]string{RefLocal, k, ""}
 				nodes[addr] = &node{
 					addr:  addr,
 					edges: func() []hcl.Traversal { return hclsyntax.Variables(v.Expr) },
@@ -160,10 +160,10 @@ func (s *State) evalReferences(ctx *hcl.EvalContext, body *hclsyntax.Body) error
 		for _, e := range n.edges() {
 			var addr [3]string
 			switch root := e.RootName(); {
-			case root == localRef && len(e) == 2:
-				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 root == RefLocal && len(e) == 2:
+				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 s.config.initblk[root] != nil && len(e) == 1:
 				addr = [3]string{root, "", ""}
 			}
@@ -181,9 +181,9 @@ func (s *State) evalReferences(ctx *hcl.EvalContext, body *hclsyntax.Body) error
 			return err
 		}
 		switch n.addr[0] {
-		case dataBlock:
+		case RefData:
 			data := make(map[string]cty.Value)
-			if vv, ok := ctx.Variables[dataBlock]; ok {
+			if vv, ok := ctx.Variables[RefData]; ok {
 				data = vv.AsValueMap()
 			}
 			src := make(map[string]cty.Value)
@@ -192,14 +192,14 @@ func (s *State) evalReferences(ctx *hcl.EvalContext, body *hclsyntax.Body) error
 			}
 			src[n.addr[2]] = v
 			data[n.addr[1]] = cty.ObjectVal(src)
-			ctx.Variables[dataBlock] = cty.ObjectVal(data)
-		case localRef:
+			ctx.Variables[RefData] = cty.ObjectVal(data)
+		case RefLocal:
 			locals := make(map[string]cty.Value)
-			if vv, ok := ctx.Variables[localRef]; ok {
+			if vv, ok := ctx.Variables[RefLocal]; ok {
 				locals = vv.AsValueMap()
 			}
 			locals[n.addr[1]] = v
-			ctx.Variables[localRef] = cty.ObjectVal(locals)
+			ctx.Variables[RefLocal] = cty.ObjectVal(locals)
 		default:
 			ctx.Variables[n.addr[0]] = v
 		}
@@ -221,15 +221,15 @@ func (s *State) evalReferences(ctx *hcl.EvalContext, body *hclsyntax.Body) error
 	return nil
 }
 
-func mergeCtxVar(ctx *hcl.EvalContext, vals map[string]cty.Value) {
-	v, ok := ctx.Variables[varRef]
+func mergeCtxVar(ctx *hcl.EvalContext, values map[string]cty.Value) {
+	v, ok := ctx.Variables[RefVar]
 	if ok {
 		v.ForEachElement(func(key cty.Value, val cty.Value) (stop bool) {
-			vals[key.AsString()] = val
+			values[key.AsString()] = val
 			return false
 		})
 	}
-	ctx.Variables[varRef] = cty.ObjectVal(vals)
+	ctx.Variables[RefVar] = cty.ObjectVal(values)
 }
 
 func setBlockVars(ctx *hcl.EvalContext, b *hclsyntax.Body) (*hcl.EvalContext, error) {
@@ -374,15 +374,60 @@ var (
 
 // Built-in blocks.
 const (
-	varBlock    = "variable"
-	dataBlock   = "data"
-	localsBlock = "locals"
-	forEachAttr = "for_each"
-	eachRef     = "each"
-	varRef      = "var"
-	localRef    = "local"
+	BlockData     = "data"
+	BlockLocals   = "locals"
+	BlockVariable = "variable"
+	RefData       = "data"
+	RefVar        = "var"
+	RefLocal      = "local"
+	AttrName      = "name"
+	forEachAttr   = "for_each"
+	eachRef       = "each"
 )
 
+// Variables represents the dynamic variables used in a body.
+type Variables struct {
+	Var, Local map[string]bool
+	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.
@@ -393,7 +438,7 @@ func defRegistry(b *hclsyntax.Body) *blockDef {
 	}
 	for _, blk := range b.Blocks {
 		// variable definition blocks are available in the HCL source but not reachable by reference.
-		if blk.Type == varBlock {
+		if blk.Type == BlockVariable {
 			continue
 		}
 		reg.child(extractDef(blk, reg))
diff --git a/schemahcl/hcl.go b/schemahcl/hcl.go
index 66595428..098110b1 100644
--- a/schemahcl/hcl.go
+++ b/schemahcl/hcl.go
@@ -90,7 +90,7 @@ func (s *State) Eval(parsed *hclparse.Parser, v any, input map[string]cty.Value)
 		for _, b := range body.Blocks {
 			switch {
 			// Variable blocks are not reachable by reference.
-			case b.Type == varBlock:
+			case b.Type == BlockVariable:
 				continue
 			// Semi-evaluate blocks with the for_each meta argument.
 			case b.Body != nil && b.Body.Attributes[forEachAttr] != nil:
@@ -222,7 +222,7 @@ func (s *State) resource(ctx *hcl.EvalContext, file *hcl.File) (*Resource, error
 	}
 	for _, blk := range body.Blocks {
 		// variable blocks may be included in the document but are skipped in unmarshaling.
-		if blk.Type == varBlock {
+		if blk.Type == BlockVariable {
 			continue
 		}
 		ctx, err := setBlockVars(ctx.NewChild(), blk.Body)


Test patch:
diff --git a/cmd/atlas/internal/cmdapi/project_test.go b/cmd/atlas/internal/cmdapi/project_test.go
index d9e07b47..4db0f228 100644
--- a/cmd/atlas/internal/cmdapi/project_test.go
+++ b/cmd/atlas/internal/cmdapi/project_test.go
@@ -198,6 +198,53 @@ env {
 	require.Equal(t, "env: local", envs[0].Format.Schema.Apply)
 }
 
+func TestPartialParse(t *testing.T) {
+	h := `
+data "remote_dir" "ignored" {
+  name = "ignored"
+}
+
+locals {
+  a = "b"
+  ignored = data.remote_dir.ignored.url
+}
+
+env {
+  name = atlas.env
+  log {
+    schema {
+      diff  = local.a
+      apply = "env: ${atlas.env}"
+    }
+  }
+  a = local.a
+}
+
+env "dev" {
+  a = local.a
+}`
+	path := filepath.Join(t.TempDir(), "atlas.hcl")
+	err := os.WriteFile(path, []byte(h), 0600)
+	require.NoError(t, err)
+	GlobalFlags.ConfigURL = "file://" + path
+	_, envs, err := EnvByName("unnamed")
+	require.NoError(t, err)
+	require.Len(t, envs, 1)
+	require.Equal(t, "unnamed", envs[0].Name)
+	require.Equal(t, "b", envs[0].Format.Schema.Diff)
+	require.Equal(t, "env: unnamed", envs[0].Format.Schema.Apply)
+	_, envs, err = EnvByName("dev")
+	require.NoError(t, err)
+	require.Len(t, envs, 2)
+	for _, e := range envs {
+		attr, ok := e.Extra.Attr("a")
+		require.True(t, ok)
+		v, err := attr.String()
+		require.NoError(t, err)
+		require.Equal(t, "b", v)
+	}
+}
+
 func TestDiff_Options(t *testing.T) {
 	d := &Diff{}
 	require.Len(t, d.Options(), 1)
diff --git a/schemahcl/context_test.go b/schemahcl/context_test.go
index 6cf0df36..f174f85f 100644
--- a/schemahcl/context_test.go
+++ b/schemahcl/context_test.go
@@ -8,6 +8,8 @@ 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"
 )
@@ -612,3 +614,77 @@ 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)
+	}
+}