Repository: ariga/atlas
Pull request #707 (2022-04-14 12:05:05)

schema/schemaspec/schemahcl: handle input variables


Patch:
diff --git a/schema/schemaspec/schemahcl/context.go b/schema/schemaspec/schemahcl/context.go
index 621c46d1..6cf7e853 100644
--- a/schema/schemaspec/schemahcl/context.go
+++ b/schema/schemaspec/schemahcl/context.go
@@ -1,3 +1,7 @@
+// Copyright 2021-present The Atlas Authors. All rights reserved.
+// This source code is licensed under the Apache 2.0 license found
+// in the LICENSE file in the root directory of this source tree.
+
 package schemahcl
 
 import (
@@ -6,7 +10,6 @@ import (
 	"strconv"
 
 	"ariga.io/atlas/schema/schemaspec"
-
 	"github.com/hashicorp/hcl/v2"
 	"github.com/hashicorp/hcl/v2/gohcl"
 	"github.com/hashicorp/hcl/v2/hclsyntax"
@@ -28,7 +31,10 @@ import (
 //		title = "{greeting.hebrew.word}, world!"
 //	}
 //
-func evalCtx(ctx *hcl.EvalContext, f *hcl.File) (*hcl.EvalContext, error) {
+func (s *State) evalCtx(ctx *hcl.EvalContext, f *hcl.File) (*hcl.EvalContext, error) {
+	if err := s.setInputVals(ctx, f.Body); err != nil {
+		return nil, err
+	}
 	c := &container{}
 	if diag := gohcl.DecodeBody(f.Body, ctx, c); diag.HasErrors() {
 		return nil, diag
@@ -40,6 +46,62 @@ func evalCtx(ctx *hcl.EvalContext, f *hcl.File) (*hcl.EvalContext, error) {
 	return setBlockVars(ctx, b)
 }
 
+// setInputVals sets the input values into the evaluation context. HCL documents can define
+// input variables in the document body by defining "variable" blocks:
+//
+// 	variable "name" {
+// 	  type = string // also supported: int, bool
+// 	  default = "rotemtam"
+// 	}
+//
+// When unmarshaling a document, input values may be provided by using the WithInputValues
+// option upon building the State:
+//
+// 	state := New(
+// 		WithInputValues(map[string]interface{}{
+// 			"name": "rotemtam",
+// 			"int":  42,
+// 			"bool": true,
+// 		}),
+// 	)
+func (s *State) setInputVals(ctx *hcl.EvalContext, body hcl.Body) error {
+	var c struct {
+		Vars []struct {
+			Name    string    `hcl:",label"`
+			Type    cty.Value `hcl:"type"`
+			Default cty.Value `hcl:"default,optional"`
+		} `hcl:"variable,block"`
+		Remain hcl.Body `hcl:",remain"`
+	}
+	nctx := ctx.NewChild()
+	nctx.Variables = map[string]cty.Value{
+		"string": capsuleTypeVal("string"),
+		"int":    capsuleTypeVal("int"),
+		"bool":   capsuleTypeVal("bool"),
+	}
+	if diag := gohcl.DecodeBody(body, nctx, &c); diag.HasErrors() {
+		return diag
+	}
+	ctxVars := make(map[string]cty.Value)
+	for _, v := range c.Vars {
+		inputVal, ok := s.config.inputValues[v.Name]
+		if ok {
+			ctxVars[v.Name] = inputVal
+			continue
+		}
+		if v.Default == cty.NilVal {
+			return fmt.Errorf("input value %q expected but missing", v.Name)
+		}
+		ctxVars[v.Name] = v.Default
+	}
+	ctx.Variables["var"] = cty.ObjectVal(ctxVars)
+	return nil
+}
+
+func capsuleTypeVal(t string) cty.Value {
+	return cty.CapsuleVal(ctyTypeSpec, &schemaspec.Type{T: t})
+}
+
 func setBlockVars(ctx *hcl.EvalContext, b *hclsyntax.Body) (*hcl.EvalContext, error) {
 	defs := defRegistry(b)
 	vars, err := blockVars(b, "", defs)
@@ -181,6 +243,9 @@ var (
 	ctyRawExpr  = cty.Capsule("raw_expr", reflect.TypeOf(schemaspec.RawExpr{}))
 )
 
+// varBlock is the block type for variable declarations.
+const varBlock = "variable"
+
 // 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.
@@ -190,6 +255,10 @@ func defRegistry(b *hclsyntax.Body) *blockDef {
 		children: make(map[string]*blockDef),
 	}
 	for _, blk := range b.Blocks {
+		// variable definition blocks are available in the HCL source but not reachable by reference.
+		if blk.Type == varBlock {
+			continue
+		}
 		reg.child(extractDef(blk, reg))
 	}
 	return reg
diff --git a/schema/schemaspec/schemahcl/hcl.go b/schema/schemaspec/schemahcl/hcl.go
index 71033f47..25c36e70 100644
--- a/schema/schemaspec/schemahcl/hcl.go
+++ b/schema/schemaspec/schemahcl/hcl.go
@@ -73,7 +73,7 @@ func (s *State) decode(ctx *hcl.EvalContext, body []byte) (*schemaspec.Resource,
 		return nil, fmt.Errorf("schemahcl: no HCL syntax found in body")
 	}
 	c := &container{}
-	ctx, err := evalCtx(ctx, srcHCL)
+	ctx, err := s.evalCtx(ctx, srcHCL)
 	if err != nil {
 		return nil, err
 	}
@@ -96,6 +96,10 @@ func (s *State) extract(ctx *hcl.EvalContext, remain hcl.Body) (*schemaspec.Reso
 		Attrs: attrs,
 	}
 	for _, blk := range body.Blocks {
+		// variable blocks may be included in the document but are skipped in unmarshaling.
+		if blk.Type == varBlock {
+			continue
+		}
 		ctx, err := setBlockVars(ctx.NewChild(), blk.Body)
 		if err != nil {
 			return nil, err
diff --git a/schema/schemaspec/schemahcl/opts.go b/schema/schemaspec/schemahcl/opts.go
index 30e3d5e5..b9dfeb9b 100644
--- a/schema/schemaspec/schemahcl/opts.go
+++ b/schema/schemaspec/schemahcl/opts.go
@@ -14,9 +14,10 @@ import (
 type (
 	// Config configures an unmarshaling.
 	Config struct {
-		types    []*schemaspec.TypeSpec
-		newCtx   func() *hcl.EvalContext
-		pathVars map[string]map[string]cty.Value
+		types       []*schemaspec.TypeSpec
+		newCtx      func() *hcl.EvalContext
+		pathVars    map[string]map[string]cty.Value
+		inputValues map[string]cty.Value
 	}
 
 	// Option configures a Config.
@@ -94,6 +95,26 @@ func WithTypes(typeSpecs []*schemaspec.TypeSpec) Option {
 	}
 }
 
+// WithInputValues configures the input values available to the Unmarshaler.
+func WithInputValues(vals map[string]interface{}) Option {
+	inputs := make(map[string]cty.Value)
+	for k, v := range vals {
+		switch v := v.(type) {
+		case string:
+			inputs[k] = cty.StringVal(v)
+		case bool:
+			inputs[k] = cty.BoolVal(v)
+		case int:
+			inputs[k] = cty.NumberIntVal(int64(v))
+		default:
+			panic("unsupported type for input value " + k)
+		}
+	}
+	return func(config *Config) {
+		config.inputValues = inputs
+	}
+}
+
 func rawExprImpl() function.Function {
 	return function.New(&function.Spec{
 		Params: []function.Parameter{
diff --git a/sql/internal/specutil/vars.go b/sql/internal/specutil/vars.go
deleted file mode 100644
index 30ea03c5..00000000
--- a/sql/internal/specutil/vars.go
+++ /dev/null
@@ -1,44 +0,0 @@
-// Copyright 2021-present The Atlas Authors. All rights reserved.
-// This source code is licensed under the Apache 2.0 license found
-// in the LICENSE file in the root directory of this source tree.
-
-package specutil
-
-import (
-	"ariga.io/atlas/schema/schemaspec"
-	"ariga.io/atlas/schema/schemaspec/schemahcl"
-)
-
-// InputVar is an input variable to an Atlas DDL document.
-type InputVar struct {
-	Name    string                   `spec:",name"`
-	Type    *schemaspec.Type         `spec:"type"`
-	Default *schemaspec.LiteralValue `spec:"default"`
-}
-
-// ExtractVarsHCL extracts the variable definitions from an Atlas DDL HCL document.
-func ExtractVarsHCL(body []byte) ([]*InputVar, error) {
-	var d struct {
-		Vars []*InputVar `spec:"variable"`
-	}
-	if err := hclState.UnmarshalSpec(body, &d); err != nil {
-		return nil, err
-	}
-	return d.Vars, nil
-}
-
-var (
-	hclState = schemahcl.New(schemahcl.WithTypes(
-		schemahcl.NewRegistry(
-			schemahcl.WithSpecs(
-				schemahcl.TypeSpec("int"),
-				schemahcl.TypeSpec("bool"),
-				schemahcl.TypeSpec("float"),
-			),
-		).Specs(),
-	))
-)
-
-func init() {
-	schemaspec.Register("variable", &InputVar{})
-}


Test patch:
diff --git a/schema/schemaspec/schemahcl/context_test.go b/schema/schemaspec/schemahcl/context_test.go
index a1ddd08c..effd52f4 100644
--- a/schema/schemaspec/schemahcl/context_test.go
+++ b/schema/schemaspec/schemahcl/context_test.go
@@ -1,3 +1,7 @@
+// Copyright 2021-present The Atlas Authors. All rights reserved.
+// This source code is licensed under the Apache 2.0 license found
+// in the LICENSE file in the root directory of this source tree.
+
 package schemahcl
 
 import (
@@ -461,3 +465,49 @@ r = user.atlas.cli
 	require.EqualValues(t, "v0.3.9", test.V)
 	require.EqualValues(t, "$user.atlas.cli", test.R.V)
 }
+
+func TestInputValues(t *testing.T) {
+	h := `
+variable "name" {
+  type = string
+}
+
+variable "default" {
+  type = string
+  default = "hello"
+}
+
+variable "int" {
+  type = int
+}
+
+variable "bool" {
+  type = bool
+}
+
+
+name = var.name
+default = var.default
+int = var.int
+bool = var.bool
+`
+	state := New(
+		WithInputValues(map[string]interface{}{
+			"name": "rotemtam",
+			"int":  42,
+			"bool": true,
+		}),
+	)
+	var test struct {
+		Name    string `spec:"name"`
+		Default string `spec:"default"`
+		Int     int    `spec:"int"`
+		Bool    bool   `spec:"bool"`
+	}
+	err := state.UnmarshalSpec([]byte(h), &test)
+	require.NoError(t, err)
+	require.EqualValues(t, "rotemtam", test.Name)
+	require.EqualValues(t, "hello", test.Default)
+	require.EqualValues(t, 42, test.Int)
+	require.EqualValues(t, true, test.Bool)
+}
diff --git a/sql/internal/specutil/vars_test.go b/sql/internal/specutil/vars_test.go
deleted file mode 100644
index a8a91d07..00000000
--- a/sql/internal/specutil/vars_test.go
+++ /dev/null
@@ -1,62 +0,0 @@
-// Copyright 2021-present The Atlas Authors. All rights reserved.
-// This source code is licensed under the Apache 2.0 license found
-// in the LICENSE file in the root directory of this source tree.
-
-package specutil
-
-import (
-	"testing"
-
-	"ariga.io/atlas/schema/schemaspec"
-	"github.com/stretchr/testify/require"
-)
-
-func TestVars(t *testing.T) {
-	h := `variable "age" {
-	type = int
-	default = 42
-}
-
-variable "happy" {
-	type = bool
-	default = true
-}
-
-variable "price" {
-	type = float
-	default = 3.14
-}
-`
-	vars, err := ExtractVarsHCL([]byte(h))
-	require.NoError(t, err)
-	require.Len(t, vars, 3)
-	require.EqualValues(t, []*InputVar{
-		{
-			Name: "age",
-			Type: &schemaspec.Type{
-				T: "int",
-			},
-			Default: &schemaspec.LiteralValue{
-				V: "42",
-			},
-		},
-		{
-			Name: "happy",
-			Type: &schemaspec.Type{
-				T: "bool",
-			},
-			Default: &schemaspec.LiteralValue{
-				V: "true",
-			},
-		},
-		{
-			Name: "price",
-			Type: &schemaspec.Type{
-				T: "float",
-			},
-			Default: &schemaspec.LiteralValue{
-				V: "3.14",
-			},
-		},
-	}, vars)
-}