Repository: ariga/atlas
Pull request #282 (2021-12-05 19:44:25)

schema/schemaspec/schemahcl: support for types as identifiers in hcl


Patch:
diff --git a/schema/schemaspec/schemahcl/context.go b/schema/schemaspec/schemahcl/context.go
index 56c8d8df..c3af557a 100644
--- a/schema/schemaspec/schemahcl/context.go
+++ b/schema/schemaspec/schemahcl/context.go
@@ -131,26 +131,29 @@ func attrMap(attrs hclsyntax.Attributes) map[string]cty.Value {
 }
 
 // ctySchemaLit is a cty.Capsule type the encapsulates a schemaspec.LiteralValue.
-var ctySchemaLit = cty.CapsuleWithOps("lit", reflect.TypeOf(schemaspec.LiteralValue{}), &cty.CapsuleOps{
-	// ConversionFrom facilitates reading the encapsulated type as a string, as is needed, for example,
-	// when interpolating it in a string expression.
-	ConversionFrom: func(src cty.Type) func(interface{}, cty.Path) (cty.Value, error) {
-		if src != cty.String {
-			return nil
-		}
-		return func(i interface{}, path cty.Path) (cty.Value, error) {
-			lit, ok := i.(*schemaspec.LiteralValue)
-			if !ok {
-				return cty.Value{}, fmt.Errorf("schemahcl: expected *schemaspec.LiteralValue got %T", i)
+var (
+	ctySchemaLit = cty.CapsuleWithOps("lit", reflect.TypeOf(schemaspec.LiteralValue{}), &cty.CapsuleOps{
+		// ConversionFrom facilitates reading the encapsulated type as a string, as is needed, for example,
+		// when interpolating it in a string expression.
+		ConversionFrom: func(src cty.Type) func(interface{}, cty.Path) (cty.Value, error) {
+			if src != cty.String {
+				return nil
 			}
-			uq, err := strconv.Unquote(lit.V)
-			if err != nil {
-				return cty.StringVal(lit.V), nil
+			return func(i interface{}, path cty.Path) (cty.Value, error) {
+				lit, ok := i.(*schemaspec.LiteralValue)
+				if !ok {
+					return cty.Value{}, fmt.Errorf("schemahcl: expected *schemaspec.LiteralValue got %T", i)
+				}
+				uq, err := strconv.Unquote(lit.V)
+				if err != nil {
+					return cty.StringVal(lit.V), nil
+				}
+				return cty.StringVal(uq), nil
 			}
-			return cty.StringVal(uq), nil
-		}
-	},
-})
+		},
+	})
+	ctySchemaType = cty.Capsule("type", reflect.TypeOf(schemaspec.Type{}))
+)
 
 // 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
diff --git a/schema/schemaspec/schemahcl/hcl.go b/schema/schemaspec/schemahcl/hcl.go
index 3adfe6ba..d9b23153 100644
--- a/schema/schemaspec/schemahcl/hcl.go
+++ b/schema/schemaspec/schemahcl/hcl.go
@@ -111,6 +111,8 @@ 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:
+			at.V = value.EncapsulatedValue().(*schemaspec.Type)
 		case value.Type().IsTupleType():
 			at.V, err = extractListValue(value)
 		default:
diff --git a/schema/schemaspec/schemahcl/opts.go b/schema/schemaspec/schemahcl/opts.go
index dfdd7d13..0c806a51 100644
--- a/schema/schemaspec/schemahcl/opts.go
+++ b/schema/schemaspec/schemahcl/opts.go
@@ -1,6 +1,7 @@
 package schemahcl
 
 import (
+	"ariga.io/atlas/schema/schemaspec"
 	"github.com/hashicorp/hcl/v2"
 	"github.com/zclconf/go-cty/cty"
 )
@@ -34,3 +35,12 @@ func EvalContext(ctx *hcl.EvalContext) Option {
 		config.ctx = ctx
 	}
 }
+
+// WithTypes configures the list of given types as identifiers in the unmarshaling context.
+func WithTypes(types []*schemaspec.Type) Option {
+	return func(config *Config) {
+		for _, t := range types {
+			config.ctx.Variables[t.Name] = cty.CapsuleVal(ctySchemaType, t)
+		}
+	}
+}
diff --git a/schema/schemaspec/spec.go b/schema/schemaspec/spec.go
index d8fc0abe..9b3a1995 100644
--- a/schema/schemaspec/spec.go
+++ b/schema/schemaspec/spec.go
@@ -63,6 +63,12 @@ type (
 	// UnmarshalerFunc is the function type that is implemented by the UnmarshalSpec
 	// method of the Unmarshaler interface.
 	UnmarshalerFunc func([]byte, interface{}) error
+
+	// Type represents the type of the field in a schema.
+	Type struct {
+		Name string
+		T    string
+	}
 )
 
 // Int returns an integer from the Value of the Attr. If The value is not a LiteralValue or the value
@@ -214,6 +220,7 @@ func BoolVal(v Value) (bool, error) {
 func (*LiteralValue) val() {}
 func (*ListValue) val()    {}
 func (*Ref) val()          {}
+func (*Type) val()         {}
 
 var (
 	_ Unmarshaler = UnmarshalerFunc(nil)


Test patch:
diff --git a/internal/integration/mysql_test.go b/internal/integration/mysql_test.go
index 54df771c..d3df78d5 100644
--- a/internal/integration/mysql_test.go
+++ b/internal/integration/mysql_test.go
@@ -525,6 +525,181 @@ func TestMySQL_CLI(t *testing.T) {
 	})
 }
 
+func TestMySQL_Sanity(t *testing.T) {
+	n := "atlas_types_sanity"
+	ddl := `
+create table atlas_types_sanity
+(
+    tBit       bit(10)       default b'100'            null,
+    tInt       int(10)       default 4             not null,
+    tTinyInt   tinyint(10)   default 8                 null,
+    tSmallInt  smallint(10)  default 2                 null,
+    tMediumInt mediumint(10) default 11                null,
+    tBigInt    bigint(10)    default 4                 null,
+    tDecimal   decimal       default 4                 null,
+    tNumeric   numeric       default 4             not null,
+    tFloat     float(10, 0)  default 4                 null,
+    tDouble    double(10, 0) default 4                 null,
+    tReal      double(10, 0) default 4                 null,
+    tTimestamp timestamp     default CURRENT_TIMESTAMP null,
+    tDate      date                                    null,
+    tTime      time                                    null,
+    tDateTime  datetime                                null,
+    tYear      year                                    null,
+    tVarchar   varchar(10)   default 'Titan'           null,
+    tChar      char(25)      default 'Olimpia'         not null,
+    tVarBinary varbinary(30) default 'Titan'           null,
+    tBinary    binary(5)     default 'Titan'           null
+) CHARSET = latin1 COLLATE latin1_swedish_ci;
+`
+	myRun(t, func(t *myTest) {
+		t.dropTables(n)
+		_, err := t.db.Exec(ddl)
+		require.NoError(t, err)
+		realm := t.loadRealm()
+		require.Len(t, realm.Schemas, 1)
+		ts, ok := realm.Schemas[0].Table(n)
+		require.True(t, ok)
+		expected := schema.Table{
+			Name: n,
+			Attrs: []schema.Attr{
+				&schema.Charset{V: "latin1"},
+				&schema.Collation{V: "latin1_swedish_ci"},
+			},
+			Schema: realm.Schemas[0],
+			Columns: []*schema.Column{
+				{
+					Name:    "tBit",
+					Type:    &schema.ColumnType{Type: &mysql.BitType{T: "bit"}, Raw: "bit(10)", Null: true},
+					Default: &schema.RawExpr{X: "b'100'"},
+				},
+				{
+					Name: "tInt",
+					Type: &schema.ColumnType{Type: &schema.IntegerType{T: "int", Unsigned: false},
+						Raw: t.valueByVersion("int(10)", "int(10)", "int"), Null: false},
+					Default: &schema.RawExpr{X: "4"},
+				},
+				{
+					Name: "tTinyInt",
+					Type: &schema.ColumnType{Type: &schema.IntegerType{T: "tinyint", Unsigned: false},
+						Raw: t.valueByVersion("tinyint(10)", "tinyint(10)", "tinyint"), Null: true},
+					Default: &schema.RawExpr{X: "8"},
+				},
+				{
+					Name: "tSmallInt",
+					Type: &schema.ColumnType{Type: &schema.IntegerType{T: "smallint", Unsigned: false},
+						Raw: t.valueByVersion("smallint(10)", "smallint(10)", "smallint"), Null: true},
+					Default: &schema.RawExpr{X: "2"},
+				},
+				{
+					Name: "tMediumInt",
+					Type: &schema.ColumnType{Type: &schema.IntegerType{T: "mediumint", Unsigned: false},
+						Raw: t.valueByVersion("mediumint(10)", "mediumint(10)", "mediumint"), Null: true},
+					Default: &schema.RawExpr{X: "11"},
+				},
+				{
+					Name: "tBigInt",
+					Type: &schema.ColumnType{Type: &schema.IntegerType{T: "bigint", Unsigned: false},
+						Raw: t.valueByVersion("bigint(10)", "bigint(10)", "bigint"), Null: true},
+					Default: &schema.RawExpr{X: "4"},
+				},
+				{
+					Name: "tDecimal",
+					Type: &schema.ColumnType{Type: &schema.DecimalType{T: "decimal", Precision: 10},
+						Raw: "decimal(10,0)", Null: true},
+					Default: &schema.RawExpr{X: "4"},
+				},
+				{
+					Name: "tNumeric",
+					Type: &schema.ColumnType{Type: &schema.DecimalType{T: "decimal", Precision: 10},
+						Raw: "decimal(10,0)", Null: false},
+					Default: &schema.RawExpr{X: "4"},
+				},
+				{
+					Name: "tFloat",
+					Type: &schema.ColumnType{Type: &schema.FloatType{T: "float", Precision: 10},
+						Raw: "float(10,0)", Null: true},
+					Default: &schema.RawExpr{X: "4"},
+				},
+				{
+					Name: "tDouble",
+					Type: &schema.ColumnType{Type: &schema.FloatType{T: "double", Precision: 10},
+						Raw: "double(10,0)", Null: true},
+					Default: &schema.RawExpr{X: "4"},
+				},
+				{
+					Name: "tReal",
+					Type: &schema.ColumnType{Type: &schema.FloatType{T: "double", Precision: 10},
+						Raw: "double(10,0)", Null: true},
+					Default: &schema.RawExpr{X: "4"},
+				},
+				{
+					Name: "tTimestamp",
+					Type: &schema.ColumnType{Type: &schema.TimeType{T: "timestamp"},
+						Raw: "timestamp", Null: true},
+					Default: &schema.RawExpr{X: "CURRENT_TIMESTAMP"},
+				},
+				{
+					Name: "tDate",
+					Type: &schema.ColumnType{Type: &schema.TimeType{T: "date"},
+						Raw: "date", Null: true},
+				},
+				{
+					Name: "tTime",
+					Type: &schema.ColumnType{Type: &schema.TimeType{T: "time"},
+						Raw: "time", Null: true},
+				},
+				{
+					Name: "tDateTime",
+					Type: &schema.ColumnType{Type: &schema.TimeType{T: "datetime"},
+						Raw: "datetime", Null: true},
+				},
+				{
+					Name: "tYear",
+					Type: &schema.ColumnType{Type: &schema.TimeType{T: "year"},
+						Raw: t.valueByVersion("year(4)", "year(4)", "year"), Null: true},
+				},
+				{
+					Name: "tVarchar",
+					Type: &schema.ColumnType{Type: &schema.StringType{T: "varchar", Size: 10},
+						Raw: "varchar(10)", Null: true},
+					Default: &schema.RawExpr{X: "Titan"},
+					Attrs: []schema.Attr{
+						&schema.Charset{V: "latin1"},
+						&schema.Collation{V: "latin1_swedish_ci"},
+					},
+				},
+				{
+					Name: "tChar",
+					Type: &schema.ColumnType{Type: &schema.StringType{T: "char", Size: 25},
+						Raw: "char(25)", Null: false},
+					Default: &schema.RawExpr{X: "Olimpia"},
+					Attrs: []schema.Attr{
+						&schema.Charset{V: "latin1"},
+						&schema.Collation{V: "latin1_swedish_ci"},
+					},
+				},
+				{
+					Name: "tVarBinary",
+					Type: &schema.ColumnType{Type: &schema.BinaryType{T: "varbinary", Size: 30},
+						Raw: "varbinary(30)", Null: true},
+					Default: &schema.RawExpr{X: t.valueByVersion("Titan", "Titan", "0x546974616E")},
+				},
+				{
+					Name: "tBinary",
+					Type: &schema.ColumnType{Type: &schema.BinaryType{T: "binary", Size: 5},
+						Raw: "binary(5)", Null: true},
+					Default: &schema.RawExpr{X: t.valueByVersion(
+						"Titan",
+						"Titan",
+						"0x546974616E")},
+				},
+			},
+		}
+		require.EqualValues(t, &expected, ts)
+	})
+}
+
 func (t *myTest) dsn() string {
 	return fmt.Sprintf("mysql://root:pass@tcp(localhost:%d)/test", t.port)
 }
@@ -636,6 +811,20 @@ func (t *myTest) posts() *schema.Table {
 	return postsT
 }
 
+func (t *myTest) valueByVersion(v56 string, v57 string, v8 string) string {
+	switch t.version {
+	case "56":
+		return v56
+	case "57":
+		return v57
+	case "8":
+		return v8
+	default:
+		t.Fatalf("unrecognized version: %q", t.version)
+		return ""
+	}
+}
+
 func (t *myTest) loadRealm() *schema.Realm {
 	r, err := t.drv.InspectRealm(context.Background(), &schema.InspectRealmOption{
 		Schemas: []string{"test"},
diff --git a/schema/schemaspec/schemahcl/context_test.go b/schema/schemaspec/schemahcl/context_test.go
index 7a9f58f4..e41bf3f1 100644
--- a/schema/schemaspec/schemahcl/context_test.go
+++ b/schema/schemaspec/schemahcl/context_test.go
@@ -349,3 +349,28 @@ x = from_ctx
 	require.NoError(t, err)
 	require.EqualValues(t, 42, test.X)
 }
+
+func TestWithTypes(t *testing.T) {
+	f := `
+first = int
+second = bool
+`
+	unmarshaler := NewUnmarshaler(
+		WithTypes(
+			[]*schemaspec.Type{
+				{Name: "int", T: "int"},
+				{Name: "bool", T: "bool"},
+			},
+		),
+	)
+	var test struct {
+		First  *schemaspec.Type `spec:"first"`
+		Second *schemaspec.Type `spec:"second"`
+	}
+	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)
+}