Repository: ariga/atlas
Pull request #3463 (2025-04-15 20:20:45)

cmd/atlas/internal/cmdapi: partial eval should detect any use of atlas.env


Patch:
diff --git a/cmd/atlas/internal/cmdapi/project.go b/cmd/atlas/internal/cmdapi/project.go
index 9b9b1233..3e4394ea 100644
--- a/cmd/atlas/internal/cmdapi/project.go
+++ b/cmd/atlas/internal/cmdapi/project.go
@@ -702,8 +702,11 @@ func partialParse(path, env string) (*hclparse.Parser, error) {
 				labeled = append(labeled, b)
 			// Dynamic env selection.
 			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 {
+				x := b.Body.Attributes[schemahcl.AttrName].Expr
+				if x != nil && schemahcl.UseTraversal(x, hcl.Traversal{
+					hcl.TraverseRoot{Name: refAtlas},
+					hcl.TraverseAttr{Name: blockEnv},
+				}) {
 					nonlabeled = append(nonlabeled, b)
 				}
 			}
diff --git a/schemahcl/schemahcl.go b/schemahcl/schemahcl.go
index 3010d2ef..5f4a37d6 100644
--- a/schemahcl/schemahcl.go
+++ b/schemahcl/schemahcl.go
@@ -9,6 +9,7 @@ import (
 	"context"
 	"fmt"
 	"reflect"
+	"slices"
 	"sort"
 	"strconv"
 	"strings"
@@ -1209,3 +1210,72 @@ func typeFuncReqArgs(spec *TypeSpec) []*TypeAttr {
 	}
 	return args
 }
+
+// UseTraversal determines if the given attribute use the given traversal.
+func UseTraversal(x hclsyntax.Expression, tr hcl.Traversal) bool {
+	switch x := x.(type) {
+	case nil:
+		return false
+	case *hclsyntax.TemplateWrapExpr:
+		return UseTraversal(x.Wrapped, tr)
+	case *hclsyntax.TemplateJoinExpr:
+		return UseTraversal(x.Tuple, tr)
+	case *hclsyntax.TemplateExpr:
+		if slices.ContainsFunc(x.Parts, func(p hclsyntax.Expression) bool {
+			return UseTraversal(p, tr)
+		}) {
+			return true
+		}
+	case *hclsyntax.ConditionalExpr:
+		for _, x1 := range []hclsyntax.Expression{x.Condition, x.TrueResult, x.FalseResult} {
+			if UseTraversal(x1, tr) {
+				return true
+			}
+		}
+	case *hclsyntax.FunctionCallExpr:
+		if slices.ContainsFunc(x.Args, func(a hclsyntax.Expression) bool {
+			return UseTraversal(a, tr)
+		}) {
+			return true
+		}
+	case *hclsyntax.ScopeTraversalExpr:
+		if len(x.Traversal) != len(tr) {
+			return false
+		}
+		if len(x.Traversal) == 0 {
+			return true
+		}
+		r1, ok1 := x.Traversal[0].(hcl.TraverseRoot)
+		r2, ok2 := tr[0].(hcl.TraverseRoot)
+		if !ok1 || !ok2 || r1.Name != r2.Name {
+			return false
+		}
+		for i := 1; i < len(x.Traversal); i++ {
+			t1, ok1 := x.Traversal[i].(hcl.TraverseAttr)
+			t2, ok2 := tr[i].(hcl.TraverseAttr)
+			if !ok1 || !ok2 || t1.Name != t2.Name {
+				return false
+			}
+		}
+		return true
+	case *hclsyntax.RelativeTraversalExpr:
+		return UseTraversal(x.Source, tr)
+	case *hclsyntax.BinaryOpExpr:
+		return UseTraversal(x.LHS, tr) || UseTraversal(x.RHS, tr)
+	case *hclsyntax.UnaryOpExpr:
+		return UseTraversal(x.Val, tr)
+	case *hclsyntax.TupleConsExpr:
+		for _, e := range x.Exprs {
+			if UseTraversal(e, tr) {
+				return true
+			}
+		}
+	case *hclsyntax.ObjectConsExpr:
+		for _, item := range x.Items {
+			if UseTraversal(item.KeyExpr, tr) || UseTraversal(item.ValueExpr, tr) {
+				return true
+			}
+		}
+	}
+	return false
+}


Test patch:
diff --git a/cmd/atlas/internal/cmdapi/project_test.go b/cmd/atlas/internal/cmdapi/project_test.go
index 521f63c3..5633e353 100644
--- a/cmd/atlas/internal/cmdapi/project_test.go
+++ b/cmd/atlas/internal/cmdapi/project_test.go
@@ -459,6 +459,25 @@ env "dev" {
 	_, envs, err = EnvByName(&cobra.Command{}, "other", nil)
 	require.Error(t, err)
 	require.Empty(t, envs)
+
+	h = `
+env {
+  name = contains(["dev", "prod", "staging"], atlas.env) ? atlas.env : "dev"
+}
+`
+	require.NoError(t, os.WriteFile(path, []byte(h), 0600))
+	_, envs, err = EnvByName(&cobra.Command{}, "none", nil)
+	require.EqualError(t, err, `env "none" not defined in config file`)
+	require.Empty(t, envs)
+
+	h = `
+env {
+  name = contains(["dev"], atlas.env) ? atlas.env : null
+}
+`
+	require.NoError(t, os.WriteFile(path, []byte(h), 0600))
+	_, envs, err = EnvByName(&cobra.Command{}, "a8m", nil)
+	require.ErrorContains(t, err, "all envs must have names on file")
 }
 
 func TestDiff_Options(t *testing.T) {
diff --git a/schemahcl/schemahcl_test.go b/schemahcl/schemahcl_test.go
index 7f83873e..db43a175 100644
--- a/schemahcl/schemahcl_test.go
+++ b/schemahcl/schemahcl_test.go
@@ -15,6 +15,7 @@ import (
 	"testing"
 
 	"github.com/hashicorp/hcl/v2"
+	"github.com/hashicorp/hcl/v2/hclparse"
 	"github.com/hashicorp/hcl/v2/hclsyntax"
 	"github.com/stretchr/testify/require"
 	"github.com/zclconf/go-cty/cty"
@@ -1218,3 +1219,44 @@ materialized "users_view" {
 	)
 	require.NoError(t, New().EvalBytes(b, &doc, nil))
 }
+
+func TestUseTraversal(t *testing.T) {
+	for _, tt := range []struct {
+		x    string
+		want bool
+	}{
+		{
+			x:    "atlas.env",
+			want: true,
+		},
+		{
+			x: "\"atlas.env\"",
+		},
+		{
+			x:    "\"${atlas.env}\"",
+			want: true,
+		},
+		{
+			x:    "contains([], atlas.env)",
+			want: true,
+		},
+		{
+			x:    "contains([], 1) ? atlas.env : 1",
+			want: true,
+		},
+		{
+			x: "contains([], 1) ? atlas.dev : 1",
+		},
+	} {
+		pr := hclparse.NewParser()
+		f, diags := pr.ParseHCL([]byte(fmt.Sprintf("x = %s", tt.x)), "test.hcl")
+		require.False(t, diags.HasErrors())
+		x := f.Body.(*hclsyntax.Body).Attributes["x"]
+		require.False(t, diags.HasErrors())
+		got := UseTraversal(x.Expr, hcl.Traversal{
+			hcl.TraverseRoot{Name: "atlas"},
+			hcl.TraverseAttr{Name: "env"},
+		})
+		require.Equal(t, tt.want, got)
+	}
+}