Repository: ariga/atlas
Pull request #288 (2021-12-07 07:08:01)

schema/schemaspec/schemahcl: support marshaling types from identifiers and functions


Patch:
diff --git a/schema/schemaspec/extension.go b/schema/schemaspec/extension.go
index f584eda8..aee6b06e 100644
--- a/schema/schemaspec/extension.go
+++ b/schema/schemaspec/extension.go
@@ -373,6 +373,9 @@ func scanPtr(key string, r *Resource, field reflect.Value) error {
 		attr.V = &LiteralValue{V: e.FieldByName("V").String()}
 	case reflect.TypeOf(Ref{}):
 		attr.V = &Ref{V: e.FieldByName("V").String()}
+	case reflect.TypeOf(Type{}):
+		t := e.Interface().(Type)
+		attr.V = &t
 	default:
 		return fmt.Errorf("schemaspec: unsupported pointer to %s", e)
 	}
diff --git a/schema/schemaspec/schemahcl/hcl.go b/schema/schemaspec/schemahcl/hcl.go
index 8005fae1..4392898a 100644
--- a/schema/schemaspec/schemahcl/hcl.go
+++ b/schema/schemaspec/schemahcl/hcl.go
@@ -2,6 +2,7 @@ package schemahcl
 
 import (
 	"bytes"
+	"errors"
 	"fmt"
 	"sort"
 	"strconv"
@@ -45,7 +46,7 @@ func (s *state) MarshalSpec(v interface{}) ([]byte, error) {
 	if err := r.Scan(v); err != nil {
 		return nil, fmt.Errorf("schemahcl: failed scanning %T to resource: %w", v, err)
 	}
-	return encode(r)
+	return s.encode(r)
 }
 
 // UnmarshalSpec implements schemaspec.Unmarshaler.
@@ -197,7 +198,7 @@ func toResource(ctx *hcl.EvalContext, block *hclsyntax.Block) (*schemaspec.Resou
 
 // encode encodes the give *schemaspec.Resource into a byte slice containing an Atlas HCL
 // document representing it.
-func encode(r *schemaspec.Resource) ([]byte, error) {
+func (s *state) encode(r *schemaspec.Resource) ([]byte, error) {
 	f := hclwrite.NewFile()
 	body := f.Body()
 	// If the resource has a Type then it is rendered as an HCL block.
@@ -206,12 +207,12 @@ func encode(r *schemaspec.Resource) ([]byte, error) {
 		body = blk.Body()
 	}
 	for _, attr := range r.Attrs {
-		if err := writeAttr(attr, body); err != nil {
+		if err := s.writeAttr(attr, body); err != nil {
 			return nil, err
 		}
 	}
 	for _, res := range r.Children {
-		if err := writeResource(res, body); err != nil {
+		if err := s.writeResource(res, body); err != nil {
 			return nil, err
 		}
 	}
@@ -220,16 +221,16 @@ func encode(r *schemaspec.Resource) ([]byte, error) {
 	return buf.Bytes(), err
 }
 
-func writeResource(b *schemaspec.Resource, body *hclwrite.Body) error {
+func (s *state) writeResource(b *schemaspec.Resource, body *hclwrite.Body) error {
 	blk := body.AppendNewBlock(b.Type, labels(b))
 	nb := blk.Body()
 	for _, attr := range b.Attrs {
-		if err := writeAttr(attr, nb); err != nil {
+		if err := s.writeAttr(attr, nb); err != nil {
 			return err
 		}
 	}
 	for _, b := range b.Children {
-		if err := writeResource(b, nb); err != nil {
+		if err := s.writeResource(b, nb); err != nil {
 			return err
 		}
 	}
@@ -244,11 +245,21 @@ func labels(r *schemaspec.Resource) []string {
 	return l
 }
 
-func writeAttr(attr *schemaspec.Attr, body *hclwrite.Body) error {
+func (s *state) writeAttr(attr *schemaspec.Attr, body *hclwrite.Body) error {
 	switch v := attr.V.(type) {
 	case *schemaspec.Ref:
 		expr := strings.ReplaceAll(v.V, "$", "")
 		body.SetAttributeRaw(attr.K, hclRawTokens(expr))
+	case *schemaspec.Type:
+		spec, ok := s.findTypeSpec(v.T)
+		if !ok {
+			return fmt.Errorf("schemahcl: type spec for %q not found", v.T)
+		}
+		s, err := hclType(spec, v)
+		if err != nil {
+			return err
+		}
+		body.SetAttributeRaw(attr.K, hclRawTokens(s))
 	case *schemaspec.LiteralValue:
 		body.SetAttributeRaw(attr.K, hclRawTokens(v.V))
 	case *schemaspec.ListValue:
@@ -271,6 +282,30 @@ func writeAttr(attr *schemaspec.Attr, body *hclwrite.Body) error {
 	return nil
 }
 
+func (s *state) findTypeSpec(t string) (*schemaspec.TypeSpec, bool) {
+	for _, v := range s.config.types {
+		if v.T == t {
+			return v, true
+		}
+	}
+	return nil, false
+}
+
+func hclType(spec *schemaspec.TypeSpec, typ *schemaspec.Type) (string, error) {
+	if len(spec.Attributes) == 0 {
+		return typ.T, nil
+	}
+	args := make([]string, 0, len(spec.Attributes))
+	for _, arg := range typ.Attributes {
+		lit, ok := arg.V.(*schemaspec.LiteralValue)
+		if !ok {
+			return "", errors.New("expecting literal value")
+		}
+		args = append(args, lit.V)
+	}
+	return fmt.Sprintf("%s(%s)", typ.T, strings.Join(args, ",")), nil
+}
+
 func hclRawTokens(s string) hclwrite.Tokens {
 	return hclwrite.Tokens{
 		&hclwrite.Token{
diff --git a/schema/schemaspec/schemahcl/opts.go b/schema/schemaspec/schemahcl/opts.go
index 79981084..7dc6d3f8 100644
--- a/schema/schemaspec/schemahcl/opts.go
+++ b/schema/schemaspec/schemahcl/opts.go
@@ -12,7 +12,8 @@ import (
 type (
 	// Config configures an unmarshaling.
 	Config struct {
-		ctx *hcl.EvalContext
+		ctx   *hcl.EvalContext
+		types []*schemaspec.TypeSpec
 	}
 
 	// Option configures a Config.
@@ -44,6 +45,7 @@ func EvalContext(ctx *hcl.EvalContext) Option {
 func WithTypes(typeSpecs []*schemaspec.TypeSpec) Option {
 	return func(config *Config) {
 		for _, typeSpec := range typeSpecs {
+			config.types = append(config.types, typeSpec)
 			if len(typeSpec.Attributes) == 0 {
 				typ := &schemaspec.Type{T: typeSpec.T}
 				config.ctx.Variables[typeSpec.Name] = cty.CapsuleVal(ctyTypeSpec, typ)


Test patch:
diff --git a/schema/schemaspec/schemahcl/context_test.go b/schema/schemaspec/schemahcl/context_test.go
index 1493f263..b7dfc4b6 100644
--- a/schema/schemaspec/schemahcl/context_test.go
+++ b/schema/schemaspec/schemahcl/context_test.go
@@ -352,10 +352,9 @@ x = from_ctx
 }
 
 func TestWithTypes(t *testing.T) {
-	f := `
-first = int
+	f := `first  = int
 second = bool
-sized = varchar(255)
+sized  = varchar(255)
 `
 	s := New(
 		WithTypes(
@@ -387,4 +386,7 @@ sized = varchar(255)
 			{K: "size", V: &schemaspec.LiteralValue{V: "255"}},
 		},
 	}, test.Varchar)
+	after, err := s.MarshalSpec(&test)
+	require.NoError(t, err)
+	require.EqualValues(t, f, string(after))
 }