Repository: ariga/atlas
Pull request #284 (2021-12-05 22:20:55)

schema/schemaspec/schemahcl: adding support for describing types using functions


Patch:
diff --git a/cmd/atlas/generate.go b/cmd/atlas/generate.go
index 80990c53..437af339 100644
--- a/cmd/atlas/generate.go
+++ b/cmd/atlas/generate.go
@@ -1,4 +1,3 @@
 package main
 
 //go:generate go run -mod=mod doc.go
-
diff --git a/schema/schemaspec/schemahcl/context.go b/schema/schemaspec/schemahcl/context.go
index c3af557a..f49700a7 100644
--- a/schema/schemaspec/schemahcl/context.go
+++ b/schema/schemaspec/schemahcl/context.go
@@ -152,7 +152,7 @@ var (
 			}
 		},
 	})
-	ctySchemaType = cty.Capsule("type", reflect.TypeOf(schemaspec.Type{}))
+	ctyTypeSpec = cty.Capsule("type", reflect.TypeOf(schemaspec.Type{}))
 )
 
 // defRegistry returns a tree of blockDef structs representing the schema of the
diff --git a/schema/schemaspec/schemahcl/hcl.go b/schema/schemaspec/schemahcl/hcl.go
index d9b23153..79b77308 100644
--- a/schema/schemaspec/schemahcl/hcl.go
+++ b/schema/schemaspec/schemahcl/hcl.go
@@ -111,7 +111,7 @@ func toAttrs(ctx *hcl.EvalContext, hclAttrs hclsyntax.Attributes) ([]*schemaspec
 		switch {
 		case isRef(value):
 			at.V = &schemaspec.Ref{V: value.GetAttr("__ref").AsString()}
-		case value.Type() == ctySchemaType:
+		case value.Type() == ctyTypeSpec:
 			at.V = value.EncapsulatedValue().(*schemaspec.Type)
 		case value.Type().IsTupleType():
 			at.V, err = extractListValue(value)
diff --git a/schema/schemaspec/schemahcl/opts.go b/schema/schemaspec/schemahcl/opts.go
index 0c806a51..842b9cc9 100644
--- a/schema/schemaspec/schemahcl/opts.go
+++ b/schema/schemaspec/schemahcl/opts.go
@@ -1,9 +1,12 @@
 package schemahcl
 
 import (
+	"reflect"
+
 	"ariga.io/atlas/schema/schemaspec"
 	"github.com/hashicorp/hcl/v2"
 	"github.com/zclconf/go-cty/cty"
+	"github.com/zclconf/go-cty/cty/function"
 )
 
 type (
@@ -21,6 +24,7 @@ func NewUnmarshaler(opts ...Option) *Unmarshaler {
 	cfg := &Config{
 		ctx: &hcl.EvalContext{
 			Variables: make(map[string]cty.Value),
+			Functions: make(map[string]function.Function),
 		},
 	}
 	for _, opt := range opts {
@@ -37,10 +41,46 @@ func EvalContext(ctx *hcl.EvalContext) Option {
 }
 
 // WithTypes configures the list of given types as identifiers in the unmarshaling context.
-func WithTypes(types []*schemaspec.Type) Option {
+func WithTypes(typeSpecs []*schemaspec.TypeSpec) Option {
 	return func(config *Config) {
-		for _, t := range types {
-			config.ctx.Variables[t.Name] = cty.CapsuleVal(ctySchemaType, t)
+		for _, typeSpec := range typeSpecs {
+			if len(typeSpec.Attributes) == 0 {
+				typ := &schemaspec.Type{T: typeSpec.T}
+				config.ctx.Variables[typeSpec.Name] = cty.CapsuleVal(ctyTypeSpec, typ)
+				continue
+			}
+			spec := &function.Spec{
+				Type: function.StaticReturnType(ctyTypeSpec),
+			}
+			for _, arg := range typeSpec.Attributes {
+				p := function.Parameter{
+					Name:      arg.Name,
+					AllowNull: !arg.Required,
+				}
+				switch arg.Kind {
+				case reflect.String:
+					p.Type = cty.String
+				case reflect.Int, reflect.Float32:
+					p.Type = cty.Number
+				case reflect.Bool:
+					p.Type = cty.Bool
+				}
+				spec.Params = append(spec.Params, p)
+				spec.Impl = func(args []cty.Value, retType cty.Type) (cty.Value, error) {
+					t := &schemaspec.Type{
+						T: typeSpec.T,
+					}
+					for i, arg := range args {
+						v, err := extractLiteralValue(arg)
+						if err != nil {
+							return cty.NilVal, err
+						}
+						t.Attributes = append(t.Attributes, &schemaspec.Attr{K: typeSpec.Attributes[i].Name, V: v})
+					}
+					return cty.CapsuleVal(ctyTypeSpec, t), nil
+				}
+			}
+			config.ctx.Functions[typeSpec.Name] = function.New(spec)
 		}
 	}
 }
diff --git a/schema/schemaspec/spec.go b/schema/schemaspec/spec.go
index 9b3a1995..3593b217 100644
--- a/schema/schemaspec/spec.go
+++ b/schema/schemaspec/spec.go
@@ -2,6 +2,7 @@ package schemaspec
 
 import (
 	"fmt"
+	"reflect"
 	"strconv"
 )
 
@@ -64,10 +65,25 @@ type (
 	// method of the Unmarshaler interface.
 	UnmarshalerFunc func([]byte, interface{}) error
 
+	// TypeSpec represents a specification for defining a Type.
+	TypeSpec struct {
+		Name       string
+		T          string
+		Attributes []*TypeAttr
+	}
+
+	// TypeAttr describes an attribute of a TypeSpec, for example `varchar` fields
+	// can have a `size` attribute.
+	TypeAttr struct {
+		Name     string
+		Kind     reflect.Kind
+		Required bool
+	}
+
 	// Type represents the type of the field in a schema.
 	Type struct {
-		Name string
-		T    string
+		T          string
+		Attributes []*Attr
 	}
 )
 


Test patch:
diff --git a/schema/schemaspec/schemahcl/context_test.go b/schema/schemaspec/schemahcl/context_test.go
index e41bf3f1..8d7cff9a 100644
--- a/schema/schemaspec/schemahcl/context_test.go
+++ b/schema/schemaspec/schemahcl/context_test.go
@@ -1,6 +1,7 @@
 package schemahcl
 
 import (
+	"reflect"
 	"testing"
 
 	"ariga.io/atlas/schema/schemaspec"
@@ -354,23 +355,36 @@ func TestWithTypes(t *testing.T) {
 	f := `
 first = int
 second = bool
+sized = varchar(255)
 `
 	unmarshaler := NewUnmarshaler(
 		WithTypes(
-			[]*schemaspec.Type{
+			[]*schemaspec.TypeSpec{
 				{Name: "int", T: "int"},
 				{Name: "bool", T: "bool"},
+				{
+					Name: "varchar",
+					T:    "varchar",
+					Attributes: []*schemaspec.TypeAttr{
+						{Name: "size", Kind: reflect.Int, Required: false},
+					},
+				},
 			},
 		),
 	)
 	var test struct {
-		First  *schemaspec.Type `spec:"first"`
-		Second *schemaspec.Type `spec:"second"`
+		First   *schemaspec.Type `spec:"first"`
+		Second  *schemaspec.Type `spec:"second"`
+		Varchar *schemaspec.Type `spec:"sized"`
 	}
 	err := unmarshaler.UnmarshalSpec([]byte(f), &test)
 	require.NoError(t, err)
-	require.EqualValues(t, "int", test.First.Name)
 	require.EqualValues(t, "int", test.First.T)
-	require.EqualValues(t, "bool", test.Second.Name)
 	require.EqualValues(t, "bool", test.Second.T)
+	require.EqualValues(t, &schemaspec.Type{
+		T: "varchar",
+		Attributes: []*schemaspec.Attr{
+			{K: "size", V: &schemaspec.LiteralValue{V: "255"}},
+		},
+	}, test.Varchar)
 }
diff --git a/sql/internal/specutil/convert_test.go b/sql/internal/specutil/convert_test.go
index 24865956..9cac1a16 100644
--- a/sql/internal/specutil/convert_test.go
+++ b/sql/internal/specutil/convert_test.go
@@ -36,7 +36,7 @@ func TestFromSpec_SchemaName(t *testing.T) {
 	})
 	require.NoError(t, err)
 	require.Equal(t, sc.Name, s.Name)
-	require.Equal(t, "$schema." + sc.Name, ta[0].Schema.V)
+	require.Equal(t, "$schema."+sc.Name, ta[0].Schema.V)
 }
 
 func TestFromForeignKey(t *testing.T) {