Repository: ariga/atlas
Pull request #304 (2021-12-13 06:20:10)

sql/mysql: add type registry
This adds the TypeRegistry for the mysql driver. It does not yet replace the older string-based processing for DDL marshal/unmarshal.


Patch:
diff --git a/schema/schemaspec/schemahcl/hcl.go b/schema/schemaspec/schemahcl/hcl.go
index 2231decf..50929a9e 100644
--- a/schema/schemaspec/schemahcl/hcl.go
+++ b/schema/schemaspec/schemahcl/hcl.go
@@ -128,7 +128,7 @@ func toAttrs(ctx *hcl.EvalContext, hclAttrs hclsyntax.Attributes) ([]*schemaspec
 		}
 		attrs = append(attrs, at)
 	}
-	// hclsyntax.Attributes is an alias for map[string]*Attribute
+	// hclsyntax.Attrs is an alias for map[string]*Attribute
 	sort.Slice(attrs, func(i, j int) bool {
 		return attrs[i].K < attrs[j].K
 	})
@@ -297,7 +297,7 @@ func hclType(spec *schemaspec.TypeSpec, typ *schemaspec.Type) (string, error) {
 	}
 	args := make([]string, 0, len(spec.Attributes))
 	for _, param := range typeFuncArgs(spec) {
-		arg, ok := findAttr(typ.Attributes, param.Name)
+		arg, ok := findAttr(typ.Attrs, param.Name)
 		if !ok {
 			continue
 		}
diff --git a/schema/schemaspec/schemahcl/opts.go b/schema/schemaspec/schemahcl/opts.go
index 13476e5a..c96357f1 100644
--- a/schema/schemaspec/schemahcl/opts.go
+++ b/schema/schemaspec/schemahcl/opts.go
@@ -87,7 +87,7 @@ func WithTypes(typeSpecs []*schemaspec.TypeSpec) Option {
 							}
 							lst.V = append(lst.V, v)
 						}
-						t.Attributes = append(t.Attributes, &schemaspec.Attr{K: spec.VarParam.Name, V: lst})
+						t.Attrs = append(t.Attrs, &schemaspec.Attr{K: spec.VarParam.Name, V: lst})
 					} else {
 						for i, arg := range args {
 							v, err := extractLiteralValue(arg)
@@ -95,7 +95,7 @@ func WithTypes(typeSpecs []*schemaspec.TypeSpec) Option {
 								return cty.NilVal, err
 							}
 							attrName := typeSpec.Attributes[i].Name
-							t.Attributes = append(t.Attributes, &schemaspec.Attr{K: attrName, V: v})
+							t.Attrs = append(t.Attrs, &schemaspec.Attr{K: attrName, V: v})
 						}
 					}
 					return cty.CapsuleVal(ctyTypeSpec, t), nil
diff --git a/schema/schemaspec/spec.go b/schema/schemaspec/spec.go
index 75a4e50a..6653611d 100644
--- a/schema/schemaspec/spec.go
+++ b/schema/schemaspec/spec.go
@@ -68,8 +68,8 @@ type (
 	// TypeSpec represents a specification for defining a Type.
 	TypeSpec struct {
 		// Name is the identifier for the type in an Atlas DDL document.
-		Name       string
-		 // T is the database identifier for the type.
+		Name string
+		// T is the database identifier for the type.
 		T          string
 		Attributes []*TypeAttr
 	}
@@ -77,15 +77,16 @@ type (
 	// TypeAttr describes an attribute of a TypeSpec, for example `varchar` fields
 	// can have a `size` attribute.
 	TypeAttr struct {
-		Name     string // Name should be a snake_case of related the schema.Type struct field.
+		// Name should be a snake_case of related the schema.Type struct field.
+		Name     string
 		Kind     reflect.Kind
 		Required bool
 	}
 
 	// Type represents the type of the field in a schema.
 	Type struct {
-		T          string
-		Attributes []*Attr
+		T     string
+		Attrs []*Attr
 	}
 )
 
diff --git a/sql/internal/specutil/types.go b/sql/internal/specutil/types.go
index 21345861..e517d6b0 100644
--- a/sql/internal/specutil/types.go
+++ b/sql/internal/specutil/types.go
@@ -3,6 +3,7 @@ package specutil
 import (
 	"errors"
 	"fmt"
+	"log"
 	"reflect"
 	"strconv"
 	"strings"
@@ -26,7 +27,7 @@ func (r *TypeRegistry) PrintType(typ *schemaspec.Type) (string, error) {
 		args        []string
 		mid, suffix string
 	)
-	for _, arg := range typ.Attributes {
+	for _, arg := range typ.Attrs {
 		// TODO(rotemtam): make this part of the TypeSpec
 		if arg.K == "unsigned" {
 			b, err := arg.Bool()
@@ -38,11 +39,24 @@ func (r *TypeRegistry) PrintType(typ *schemaspec.Type) (string, error) {
 			}
 			continue
 		}
-		lit, ok := arg.V.(*schemaspec.LiteralValue)
-		if !ok {
-			return "", errors.New("expecting literal value")
+		switch v := arg.V.(type) {
+		case *schemaspec.LiteralValue:
+			args = append(args, v.V)
+		case *schemaspec.ListValue:
+			for _, li := range v.V {
+				lit, ok := li.(*schemaspec.LiteralValue)
+				if !ok {
+					return "", fmt.Errorf("expecting literal value. got: %T", li)
+				}
+				uq, err := strconv.Unquote(lit.V)
+				if err != nil {
+					return "", fmt.Errorf("expecting list items to be quoted strings: %w", err)
+				}
+				args = append(args, "'"+uq+"'")
+			}
+		default:
+			return "", fmt.Errorf("unsupported type %T for PrintType", v)
 		}
-		args = append(args, lit.V)
 	}
 	if len(args) > 0 {
 		mid = "(" + strings.Join(args, ",") + ")"
@@ -69,6 +83,16 @@ func (r *TypeRegistry) Register(specs ...*schemaspec.TypeSpec) error {
 	return nil
 }
 
+// NewRegistry creates a new *TypeRegistry, registers the provided types and panics
+// if an error occurs.
+func NewRegistry(specs ...*schemaspec.TypeSpec) *TypeRegistry {
+	r := &TypeRegistry{}
+	if err := r.Register(specs...); err != nil {
+		log.Fatalf("failed registering types: %s", err)
+	}
+	return r
+}
+
 // FindByName searches the registry for types that have the provided name.
 func (r *TypeRegistry) FindByName(name string) (*schemaspec.TypeSpec, bool) {
 	for _, current := range r.r {
@@ -123,10 +147,10 @@ func (r *TypeRegistry) Convert(typ schema.Type) (*schemaspec.Type, error) {
 		switch attr.Kind {
 		case reflect.Int:
 			i := strconv.Itoa(int(field.Int()))
-			s.Attributes = append(s.Attributes, LitAttr(attr.Name, i))
+			s.Attrs = append(s.Attrs, LitAttr(attr.Name, i))
 		case reflect.Bool:
 			b := strconv.FormatBool(field.Bool())
-			s.Attributes = append(s.Attributes, LitAttr(attr.Name, b))
+			s.Attrs = append(s.Attrs, LitAttr(attr.Name, b))
 		case reflect.Slice:
 			lits := make([]string, 0, field.Len())
 			for i := 0; i < field.Len(); i++ {
@@ -136,7 +160,7 @@ func (r *TypeRegistry) Convert(typ schema.Type) (*schemaspec.Type, error) {
 				}
 				lits = append(lits, strconv.Quote(fi.String()))
 			}
-			s.Attributes = append(s.Attributes, ListAttr(attr.Name, lits...))
+			s.Attrs = append(s.Attrs, ListAttr(attr.Name, lits...))
 		default:
 			return nil, fmt.Errorf("specutil: unsupported attr kind %s for attribute %q of %q", attr.Kind, attr.Name, typeSpec.Name)
 		}
@@ -144,6 +168,27 @@ func (r *TypeRegistry) Convert(typ schema.Type) (*schemaspec.Type, error) {
 	return s, nil
 }
 
+// Specs returns the TypeSpecs in the registry.
+func (r *TypeRegistry) Specs() []*schemaspec.TypeSpec {
+	return r.r
+}
+
+// Type converts a *schemaspec.Type into a schema.Type.
+func (r *TypeRegistry) Type(typ *schemaspec.Type, extra []*schemaspec.Attr, parser func(string) (schema.Type, error)) (schema.Type, error) {
+	typeSpec, ok := r.Find(typ.T)
+	if !ok {
+		return nil, fmt.Errorf("specutil: typespec not found for %s", typ.T)
+	}
+	nfa := typeNonFuncArgs(typeSpec)
+	picked := pickTypeAttrs(extra, nfa)
+	typ.Attrs = appendIfNotExist(typ.Attrs, picked)
+	printType, err := r.PrintType(typ)
+	if err != nil {
+		return nil, err
+	}
+	return parser(printType)
+}
+
 // TypeSpec returns a TypeSpec with the provided name.
 func TypeSpec(name string, attrs ...*schemaspec.TypeAttr) *schemaspec.TypeSpec {
 	return &schemaspec.TypeSpec{
@@ -161,3 +206,44 @@ func SizeTypeAttr(required bool) *schemaspec.TypeAttr {
 		Required: required,
 	}
 }
+
+// typeNonFuncArgs returns the type attributes that are NOT configured via arguments to the
+// type definition, `int unsigned`.
+func typeNonFuncArgs(spec *schemaspec.TypeSpec) []*schemaspec.TypeAttr {
+	var args []*schemaspec.TypeAttr
+	for _, attr := range spec.Attributes {
+		// TODO(rotemtam): this should be defined on the TypeSpec.
+		if attr.Name == "unsigned" {
+			args = append(args, attr)
+		}
+	}
+	return args
+}
+
+// pickTypeAttrs returns the relevant Attrs matching the wanted TypeAttrs.
+func pickTypeAttrs(src []*schemaspec.Attr, wanted []*schemaspec.TypeAttr) []*schemaspec.Attr {
+	keys := make(map[string]struct{})
+	for _, w := range wanted {
+		keys[w.Name] = struct{}{}
+	}
+	var picked []*schemaspec.Attr
+	for _, attr := range src {
+		if _, ok := keys[attr.K]; ok {
+			picked = append(picked, attr)
+		}
+	}
+	return picked
+}
+
+func appendIfNotExist(base []*schemaspec.Attr, additional []*schemaspec.Attr) []*schemaspec.Attr {
+	exists := make(map[string]struct{})
+	for _, attr := range base {
+		exists[attr.K] = struct{}{}
+	}
+	for _, attr := range additional {
+		if _, ok := exists[attr.K]; !ok {
+			base = append(base, attr)
+		}
+	}
+	return base
+}
diff --git a/sql/mysql/sqlspec.go b/sql/mysql/sqlspec.go
index a969aea3..2aaf06cd 100644
--- a/sql/mysql/sqlspec.go
+++ b/sql/mysql/sqlspec.go
@@ -4,6 +4,7 @@ import (
 	"errors"
 	"fmt"
 	"math"
+	"reflect"
 	"strconv"
 	"strings"
 
@@ -436,3 +437,28 @@ func hasCollate(attr []schema.Attr, parent []schema.Attr) (string, bool) {
 	}
 	return "", false
 }
+
+// TypeRegistry contains the supported TypeSpecs for the mysql driver.
+var TypeRegistry = specutil.NewRegistry(
+	specutil.TypeSpec(tInt, unsignedTypeAttr()),
+	specutil.TypeSpec(tBigInt, unsignedTypeAttr()),
+	specutil.TypeSpec(tTinyInt, unsignedTypeAttr()),
+	specutil.TypeSpec(tSmallInt, unsignedTypeAttr()),
+	specutil.TypeSpec(tMediumInt, unsignedTypeAttr()),
+	specutil.TypeSpec(tText),
+	specutil.TypeSpec(tTinyText),
+	specutil.TypeSpec(tMediumText),
+	specutil.TypeSpec(tLongText),
+	specutil.TypeSpec(tVarchar, specutil.SizeTypeAttr(true)),
+	specutil.TypeSpec(tChar, specutil.SizeTypeAttr(true)),
+	specutil.TypeSpec(tBinary, specutil.SizeTypeAttr(true)),
+	specutil.TypeSpec(tVarBinary, specutil.SizeTypeAttr(true)),
+	specutil.TypeSpec(tEnum, &schemaspec.TypeAttr{Name: "values", Kind: reflect.Slice, Required: true}),
+)
+
+func unsignedTypeAttr() *schemaspec.TypeAttr {
+	return &schemaspec.TypeAttr{
+		Name: "unsigned",
+		Kind: reflect.Bool,
+	}
+}
diff --git a/sql/schema/schema.go b/sql/schema/schema.go
index 414181f2..cce1c5c5 100644
--- a/sql/schema/schema.go
+++ b/sql/schema/schema.go
@@ -17,7 +17,7 @@ type (
 		Name   string
 		Realm  *Realm
 		Tables []*Table
-		Attrs  []Attr // Attributes and options.
+		Attrs  []Attr // Attrs and options.
 	}
 
 	// A Table represents a table definition.
@@ -28,7 +28,7 @@ type (
 		Indexes     []*Index
 		PrimaryKey  *Index
 		ForeignKeys []*ForeignKey
-		Attrs       []Attr // Attributes, constraints and options.
+		Attrs       []Attr // Attrs, constraints and options.
 	}
 
 	// A Column represents a column definition.


Test patch:
diff --git a/schema/schemaspec/schemahcl/context_test.go b/schema/schemaspec/schemahcl/context_test.go
index d4826d2c..ca439069 100644
--- a/schema/schemaspec/schemahcl/context_test.go
+++ b/schema/schemaspec/schemahcl/context_test.go
@@ -397,13 +397,13 @@ variadic = enum("a","b","c")
 	require.EqualValues(t, "bool", test.Second.T)
 	require.EqualValues(t, &schemaspec.Type{
 		T: "varchar",
-		Attributes: []*schemaspec.Attr{
+		Attrs: []*schemaspec.Attr{
 			{K: "size", V: &schemaspec.LiteralValue{V: "255"}},
 		},
 	}, test.Varchar)
 	require.EqualValues(t, &schemaspec.Type{
 		T: "enum",
-		Attributes: []*schemaspec.Attr{
+		Attrs: []*schemaspec.Attr{
 			{
 				K: "values",
 				V: &schemaspec.ListValue{
diff --git a/sql/internal/specutil/types_test.go b/sql/internal/specutil/types_test.go
index 160e127d..4ee551c4 100644
--- a/sql/internal/specutil/types_test.go
+++ b/sql/internal/specutil/types_test.go
@@ -29,7 +29,7 @@ func TestTypePrint(t *testing.T) {
 		},
 		{
 			spec:     intSpec,
-			typ:      &schemaspec.Type{T: "int", Attributes: []*schemaspec.Attr{LitAttr("unsigned", "true")}},
+			typ:      &schemaspec.Type{T: "int", Attrs: []*schemaspec.Attr{LitAttr("unsigned", "true")}},
 			expected: "int unsigned",
 		},
 		{
@@ -40,7 +40,7 @@ func TestTypePrint(t *testing.T) {
 					{Name: "size", Kind: reflect.Int, Required: true},
 				},
 			},
-			typ:      &schemaspec.Type{T: "varchar", Attributes: []*schemaspec.Attr{LitAttr("size", "255")}},
+			typ:      &schemaspec.Type{T: "varchar", Attrs: []*schemaspec.Attr{LitAttr("size", "255")}},
 			expected: "varchar(255)",
 		},
 	} {
@@ -99,26 +99,26 @@ func TestRegistryConvert(t *testing.T) {
 	}{
 		{
 			typ:      &schema.StringType{T: "varchar", Size: 255},
-			expected: &schemaspec.Type{T: "varchar", Attributes: []*schemaspec.Attr{LitAttr("size", "255")}},
+			expected: &schemaspec.Type{T: "varchar", Attrs: []*schemaspec.Attr{LitAttr("size", "255")}},
 		},
 		{
 			typ:      &schema.IntegerType{T: "int", Unsigned: true},
-			expected: &schemaspec.Type{T: "int", Attributes: []*schemaspec.Attr{LitAttr("unsigned", "true")}},
+			expected: &schemaspec.Type{T: "int", Attrs: []*schemaspec.Attr{LitAttr("unsigned", "true")}},
 		},
 		{
 			typ:      &schema.IntegerType{T: "int", Unsigned: true},
-			expected: &schemaspec.Type{T: "int", Attributes: []*schemaspec.Attr{LitAttr("unsigned", "true")}},
+			expected: &schemaspec.Type{T: "int", Attrs: []*schemaspec.Attr{LitAttr("unsigned", "true")}},
 		},
 		{
 			typ: &schema.DecimalType{T: "decimal", Precision: 10, Scale: 2},
-			expected: &schemaspec.Type{T: "decimal", Attributes: []*schemaspec.Attr{
+			expected: &schemaspec.Type{T: "decimal", Attrs: []*schemaspec.Attr{
 				LitAttr("precision", "10"),
 				LitAttr("scale", "2"),
 			}},
 		},
 		{
 			typ: &schema.EnumType{T: "enum", Values: []string{"on", "off"}},
-			expected: &schemaspec.Type{T: "enum", Attributes: []*schemaspec.Attr{
+			expected: &schemaspec.Type{T: "enum", Attrs: []*schemaspec.Attr{
 				ListAttr("values", `"on"`, `"off"`),
 			}},
 		},
diff --git a/sql/mysql/sqlspec_test.go b/sql/mysql/sqlspec_test.go
index 62a75234..a2a1d4f5 100644
--- a/sql/mysql/sqlspec_test.go
+++ b/sql/mysql/sqlspec_test.go
@@ -5,14 +5,18 @@ import (
 	"log"
 	"testing"
 
+	"ariga.io/atlas/schema/schemaspec"
 	"ariga.io/atlas/schema/schemaspec/schemahcl"
 	"ariga.io/atlas/sql/internal/specutil"
 	"ariga.io/atlas/sql/schema"
 	"ariga.io/atlas/sql/sqlspec"
+	"github.com/hashicorp/hcl/v2/hclwrite"
 
 	"github.com/stretchr/testify/require"
 )
 
+var hclState = schemahcl.New(schemahcl.WithTypes(TypeRegistry.Specs()))
+
 func TestSQLSpec(t *testing.T) {
 	f := `
 schema "schema" {
@@ -629,3 +633,128 @@ func TestMarshalSpecColumnType(t *testing.T) {
 		})
 	}
 }
+
+func TestTypes(t *testing.T) {
+	for _, tt := range []struct {
+		typeExpr  string
+		extraAttr string
+		expected  schema.Type
+	}{
+		{
+			typeExpr: "varchar(255)",
+			expected: &schema.StringType{T: tVarchar, Size: 255},
+		},
+		{
+			typeExpr: "char(255)",
+			expected: &schema.StringType{T: tChar, Size: 255},
+		},
+		{
+			typeExpr: "binary(255)",
+			expected: &schema.BinaryType{T: tBinary, Size: 255},
+		},
+		{
+			typeExpr: "varbinary(255)",
+			expected: &schema.BinaryType{T: tVarBinary, Size: 255},
+		},
+		{
+			typeExpr: "int",
+			expected: &schema.IntegerType{T: tInt},
+		},
+		{
+			typeExpr:  "int",
+			extraAttr: "unsigned=true",
+			expected:  &schema.IntegerType{T: tInt, Unsigned: true},
+		},
+		{
+			typeExpr: "bigint",
+			expected: &schema.IntegerType{T: tBigInt},
+		},
+		{
+			typeExpr:  "bigint",
+			extraAttr: "unsigned=true",
+			expected:  &schema.IntegerType{T: tBigInt, Unsigned: true},
+		},
+		{
+			typeExpr: "tinyint",
+			expected: &schema.IntegerType{T: tTinyInt},
+		},
+		{
+			typeExpr:  "tinyint",
+			extraAttr: "unsigned=true",
+			expected:  &schema.IntegerType{T: tTinyInt, Unsigned: true},
+		},
+		{
+			typeExpr: "smallint",
+			expected: &schema.IntegerType{T: tSmallInt},
+		},
+		{
+			typeExpr:  "smallint",
+			extraAttr: "unsigned=true",
+			expected:  &schema.IntegerType{T: tSmallInt, Unsigned: true},
+		},
+		{
+			typeExpr: "mediumint",
+			expected: &schema.IntegerType{T: tMediumInt},
+		},
+		{
+			typeExpr:  "mediumint",
+			extraAttr: "unsigned=true",
+			expected:  &schema.IntegerType{T: tMediumInt, Unsigned: true},
+		},
+		{
+			typeExpr: "tinytext",
+			expected: &schema.StringType{T: tTinyText},
+		},
+		{
+			typeExpr: "mediumtext",
+			expected: &schema.StringType{T: tMediumText},
+		},
+		{
+			typeExpr: "longtext",
+			expected: &schema.StringType{T: tLongText},
+		},
+		{
+			typeExpr: "text",
+			expected: &schema.StringType{T: tText},
+		},
+		{
+			typeExpr: `enum("on","off")`,
+			expected: &schema.EnumType{Values: []string{"on", "off"}},
+		},
+	} {
+		t.Run(tt.typeExpr, func(t *testing.T) {
+			// simulates sqlspec.Column until we change its Type field.
+			type col struct {
+				Type *schemaspec.Type `spec:"type"`
+				schemaspec.DefaultExtension
+			}
+			var test struct {
+				Columns []*col `spec:"column"`
+			}
+			doc := fmt.Sprintf(`column {
+	type = %s%s
+}
+`, tt.typeExpr, lineIfSet(tt.extraAttr))
+			err := hclState.UnmarshalSpec([]byte(doc), &test)
+			require.NoError(t, err)
+			column := test.Columns[0]
+			typ, err := TypeRegistry.Type(column.Type, column.Extra.Attrs, parseRawType)
+			require.NoError(t, err)
+			require.EqualValues(t, tt.expected, typ)
+			spec, err := hclState.MarshalSpec(&test)
+			require.NoError(t, err)
+			hclEqual(t, []byte(doc), spec)
+		})
+	}
+}
+
+func hclEqual(t *testing.T, expected, actual []byte) {
+	require.EqualValues(t, string(hclwrite.Format(expected)), string(hclwrite.Format(actual)))
+}
+
+func lineIfSet(s string) string {
+	if s != "" {
+		return "\n" + s
+	}
+	return s
+}