Repository: ariga/atlas
Pull request #276 (2021-12-03 11:10:15)

sql/mysql: add support for charset attributes in sqlspec


Patch:
diff --git a/go.mod b/go.mod
index 5518cf33..31eb89e1 100644
--- a/go.mod
+++ b/go.mod
@@ -1,6 +1,6 @@
 module ariga.io/atlas
 
-go 1.16
+go 1.17
 
 require (
 	github.com/DATA-DOG/go-sqlmock v1.5.0
@@ -18,3 +18,15 @@ require (
 	golang.org/x/text v0.3.7 // indirect
 	gopkg.in/check.v1 v1.0.0-20190902080502-41f04d3bba15 // indirect
 )
+
+require (
+	github.com/agext/levenshtein v1.2.1 // indirect
+	github.com/apparentlymart/go-textseg/v13 v13.0.0 // indirect
+	github.com/chzyer/readline v0.0.0-20180603132655-2972be24d48e // indirect
+	github.com/davecgh/go-spew v1.1.1 // indirect
+	github.com/inconshreveable/mousetrap v1.0.0 // indirect
+	github.com/mitchellh/go-wordwrap v0.0.0-20150314170334-ad45545899c7 // indirect
+	github.com/pmezard/go-difflib v1.0.0 // indirect
+	github.com/spf13/pflag v1.0.5 // indirect
+	gopkg.in/yaml.v3 v3.0.0-20210107192922-496545a6307b // indirect
+)
diff --git a/internal/integration/go.mod b/internal/integration/go.mod
index 5205c2ca..7e68b1e1 100644
--- a/internal/integration/go.mod
+++ b/internal/integration/go.mod
@@ -18,6 +18,7 @@ require (
 require (
 	github.com/agext/levenshtein v1.2.3 // indirect
 	github.com/apparentlymart/go-textseg/v13 v13.0.0 // indirect
+	github.com/chzyer/readline v0.0.0-20180603132655-2972be24d48e // indirect
 	github.com/davecgh/go-spew v1.1.1 // indirect
 	github.com/google/go-cmp v0.5.6 // indirect
 	github.com/google/uuid v1.3.0 // indirect
diff --git a/sql/internal/specutil/convert.go b/sql/internal/specutil/convert.go
index 4c123333..3e6dbed5 100644
--- a/sql/internal/specutil/convert.go
+++ b/sql/internal/specutil/convert.go
@@ -21,7 +21,7 @@ type (
 	ConvertTypeFunc       func(*sqlspec.Column) (schema.Type, error)
 	ConvertPrimaryKeyFunc func(*sqlspec.PrimaryKey, *schema.Table) (*schema.Index, error)
 	ConvertIndexFunc      func(*sqlspec.Index, *schema.Table) (*schema.Index, error)
-	ColumnSpecFunc        func(*schema.Column) (*sqlspec.Column, error)
+	ColumnSpecFunc        func(*schema.Column, *schema.Table) (*sqlspec.Column, error)
 	TableSpecFunc         func(*schema.Table) (*sqlspec.Table, error)
 	PrimaryKeySpecFunc    func(index *schema.Index) (*sqlspec.PrimaryKey, error)
 	IndexSpecFunc         func(index *schema.Index) (*sqlspec.Index, error)
@@ -238,7 +238,7 @@ func FromTable(t *schema.Table, colFn ColumnSpecFunc, pkFn PrimaryKeySpecFunc, i
 		Name: t.Name,
 	}
 	for _, c := range t.Columns {
-		col, err := colFn(c)
+		col, err := colFn(c, t)
 		if err != nil {
 			return nil, err
 		}
diff --git a/sql/internal/specutil/spec.go b/sql/internal/specutil/spec.go
index 77a6b002..2572b677 100644
--- a/sql/internal/specutil/spec.go
+++ b/sql/internal/specutil/spec.go
@@ -1,6 +1,8 @@
 package specutil
 
 import (
+	"strconv"
+
 	"ariga.io/atlas/schema/schemaspec"
 	"ariga.io/atlas/sql/sqlspec"
 )
@@ -16,6 +18,14 @@ func NewCol(name, coltype string, attrs ...*schemaspec.Attr) *sqlspec.Column {
 	}
 }
 
+// StrAttr is a helper method for constructing *schemaspec.Attr of type string.
+func StrAttr(k, v string) *schemaspec.Attr {
+	return &schemaspec.Attr{
+		K: k,
+		V: &schemaspec.LiteralValue{V: strconv.Quote(v)},
+	}
+}
+
 // LitAttr is a helper method for constructing *schemaspec.Attr instances that contain literal values.
 func LitAttr(k, v string) *schemaspec.Attr {
 	return &schemaspec.Attr{
diff --git a/sql/mysql/sqlspec.go b/sql/mysql/sqlspec.go
index a5b95797..a969aea3 100644
--- a/sql/mysql/sqlspec.go
+++ b/sql/mysql/sqlspec.go
@@ -9,6 +9,7 @@ import (
 
 	"ariga.io/atlas/schema/schemaspec"
 	"ariga.io/atlas/sql/internal/specutil"
+	"ariga.io/atlas/sql/internal/sqlx"
 	"ariga.io/atlas/sql/schema"
 	"ariga.io/atlas/sql/sqlspec"
 )
@@ -24,18 +25,22 @@ func UnmarshalSpec(data []byte, unmarshaler schemaspec.Unmarshaler, v interface{
 	if err := unmarshaler.UnmarshalSpec(data, &d); err != nil {
 		return err
 	}
-	if v, ok := v.(*schema.Schema); ok {
-		if len(d.Schemas) != 1 {
-			return fmt.Errorf("mysql: expecting document to contain a single schema, got %d", len(d.Schemas))
-		}
-		conv, err := specutil.Schema(d.Schemas[0], d.Tables, convertTable)
-		if err != nil {
-			return fmt.Errorf("mysql: failed converting to *schema.Schema: %w", err)
-		}
-		*v = *conv
-		return nil
+	s, ok := v.(*schema.Schema)
+	if !ok {
+		return fmt.Errorf("mysql: failed unmarshaling spec. %T is not supported", v)
+	}
+	if len(d.Schemas) != 1 {
+		return fmt.Errorf("mysql: expecting document to contain a single schema, got %d", len(d.Schemas))
+	}
+	conv, err := specutil.Schema(d.Schemas[0], d.Tables, convertTable)
+	if err != nil {
+		return fmt.Errorf("mysql: failed converting to *schema.Schema: %w", err)
+	}
+	if err := convertCharset(d.Schemas[0], &conv.Attrs); err != nil {
+		return err
 	}
-	return fmt.Errorf("mysql: failed unmarshaling spec. %T is not supported", v)
+	*s = *conv
+	return nil
 }
 
 // MarshalSpec marshals v into an Atlas DDL document using a schemaspec.Marshaler.
@@ -61,7 +66,14 @@ func MarshalSpec(v interface{}, marshaler schemaspec.Marshaler) ([]byte, error)
 // ForeignKeySpecs into ForeignKeys, as the target tables do not necessarily exist in the schema
 // at this point. Instead, the linking is done by the convertSchema function.
 func convertTable(spec *sqlspec.Table, parent *schema.Schema) (*schema.Table, error) {
-	return specutil.Table(spec, parent, convertColumn, convertPrimaryKey, convertIndex)
+	t, err := specutil.Table(spec, parent, convertColumn, convertPrimaryKey, convertIndex)
+	if err != nil {
+		return nil, err
+	}
+	if err := convertCharset(spec, &t.Attrs); err != nil {
+		return nil, err
+	}
+	return t, err
 }
 
 // convertPrimaryKey converts a sqlspec.PrimaryKey to a schema.Index.
@@ -76,7 +88,14 @@ func convertIndex(spec *sqlspec.Index, parent *schema.Table) (*schema.Index, err
 
 // convertColumn converts a sqlspec.Column into a schema.Column.
 func convertColumn(spec *sqlspec.Column, _ *schema.Table) (*schema.Column, error) {
-	return specutil.Column(spec, convertColumnType)
+	c, err := specutil.Column(spec, convertColumnType)
+	if err != nil {
+		return nil, err
+	}
+	if err := convertCharset(spec, &c.Attrs); err != nil {
+		return nil, err
+	}
+	return c, err
 }
 
 // convertColumnType converts a sqlspec.Column into a concrete MySQL schema.Type.
@@ -235,25 +254,51 @@ func convertFloat(spec *sqlspec.Column) (schema.Type, error) {
 }
 
 // schemaSpec converts from a concrete MySQL schema to Atlas specification.
-func schemaSpec(schem *schema.Schema) (*sqlspec.Schema, []*sqlspec.Table, error) {
-	return specutil.FromSchema(schem, tableSpec)
+func schemaSpec(s *schema.Schema) (*sqlspec.Schema, []*sqlspec.Table, error) {
+	sc, t, err := specutil.FromSchema(s, tableSpec)
+	if err != nil {
+		return nil, nil, err
+	}
+	if c, ok := hasCharset(s.Attrs, nil); ok {
+		sc.Extra.Attrs = append(sc.Extra.Attrs, specutil.StrAttr("charset", c))
+	}
+	if c, ok := hasCollate(s.Attrs, nil); ok {
+		sc.Extra.Attrs = append(sc.Extra.Attrs, specutil.StrAttr("collation", c))
+	}
+	return sc, t, nil
 }
 
 // tableSpec converts from a concrete MySQL sqlspec.Table to a schema.Table.
-func tableSpec(tab *schema.Table) (*sqlspec.Table, error) {
-	return specutil.FromTable(tab, columnSpec, specutil.FromPrimaryKey, specutil.FromIndex, specutil.FromForeignKey)
+func tableSpec(t *schema.Table) (*sqlspec.Table, error) {
+	ts, err := specutil.FromTable(t, columnSpec, specutil.FromPrimaryKey, specutil.FromIndex, specutil.FromForeignKey)
+	if err != nil {
+		return nil, err
+	}
+	if c, ok := hasCharset(t.Attrs, t.Schema.Attrs); ok {
+		ts.Extra.Attrs = append(ts.Extra.Attrs, specutil.StrAttr("charset", c))
+	}
+	if c, ok := hasCollate(t.Attrs, t.Schema.Attrs); ok {
+		ts.Extra.Attrs = append(ts.Extra.Attrs, specutil.StrAttr("collation", c))
+	}
+	return ts, nil
 }
 
 // columnSpec converts from a concrete MySQL schema.Column into a sqlspec.Column.
-func columnSpec(col *schema.Column) (*sqlspec.Column, error) {
-	ct, err := columnTypeSpec(col.Type.Type)
+func columnSpec(c *schema.Column, t *schema.Table) (*sqlspec.Column, error) {
+	ct, err := columnTypeSpec(c.Type.Type)
 	if err != nil {
 		return nil, err
 	}
+	if c, ok := hasCharset(c.Attrs, t.Attrs); ok {
+		ct.Extra.Attrs = append(ct.Extra.Attrs, specutil.StrAttr("charset", c))
+	}
+	if c, ok := hasCollate(c.Attrs, t.Attrs); ok {
+		ct.Extra.Attrs = append(ct.Extra.Attrs, specutil.StrAttr("collation", c))
+	}
 	return &sqlspec.Column{
-		Name: col.Name,
+		Name: c.Name,
 		Type: ct.Type,
-		Null: col.Type.Null,
+		Null: c.Type.Null,
 		DefaultExtension: schemaspec.DefaultExtension{
 			Extra: schemaspec.Resource{Attrs: ct.DefaultExtension.Extra.Attrs},
 		},
@@ -307,8 +352,10 @@ func binarySpec(t *schema.BinaryType) (*sqlspec.Column, error) {
 func stringSpec(t *schema.StringType) (*sqlspec.Column, error) {
 	switch t.T {
 	case tVarchar, tMediumText, tLongText, tTinyText, tText, tChar:
-		s := strconv.Itoa(t.Size)
-		return specutil.NewCol("", "string", specutil.LitAttr("size", s)), nil
+		if t.Size == 0 {
+			return &sqlspec.Column{Type: t.T}, nil
+		}
+		return specutil.NewCol("", "string", specutil.LitAttr("size", strconv.Itoa(t.Size))), nil
 	}
 	return nil, fmt.Errorf("mysql: schema string failed to convert %q", t.T)
 }
@@ -345,3 +392,47 @@ func enumSpec(t *schema.EnumType) (*sqlspec.Column, error) {
 	}
 	return specutil.NewCol("", "enum", specutil.ListAttr("values", quoted...)), nil
 }
+
+// convertCharset converts spec charset/collation
+// attributes to schema element attributes.
+func convertCharset(spec interface {
+	Attr(string) (*schemaspec.Attr, bool)
+}, attrs *[]schema.Attr) error {
+	if attr, ok := spec.Attr("charset"); ok {
+		s, err := attr.String()
+		if err != nil {
+			return err
+		}
+		*attrs = append(*attrs, &schema.Charset{V: s})
+	}
+	if attr, ok := spec.Attr("collation"); ok {
+		s, err := attr.String()
+		if err != nil {
+			return err
+		}
+		*attrs = append(*attrs, &schema.Collation{V: s})
+	}
+	return nil
+}
+
+// hasCharset reports if the attribute contains the "charset" attribute,
+// and it needs to be defined explicitly on the schema. This is true, in
+// case the element charset is different from its parent charset.
+func hasCharset(attr []schema.Attr, parent []schema.Attr) (string, bool) {
+	var c, p schema.Charset
+	if sqlx.Has(attr, &c) && (parent == nil || sqlx.Has(parent, &p) && c.V != p.V) {
+		return c.V, true
+	}
+	return "", false
+}
+
+// hasCollate reports if the attribute contains the "collation" attribute,
+// and it needs to be defined explicitly on the schema. This is true, in
+// case the element collation is different from its parent collation.
+func hasCollate(attr []schema.Attr, parent []schema.Attr) (string, bool) {
+	var c, p schema.Collation
+	if sqlx.Has(attr, &c) && (parent == nil || sqlx.Has(parent, &p) && c.V != p.V) {
+		return c.V, true
+	}
+	return "", false
+}
diff --git a/sql/postgres/sqlspec.go b/sql/postgres/sqlspec.go
index 7442183c..cbf3fd09 100644
--- a/sql/postgres/sqlspec.go
+++ b/sql/postgres/sqlspec.go
@@ -226,7 +226,7 @@ func tableSpec(tab *schema.Table) (*sqlspec.Table, error) {
 }
 
 // columnSpec converts from a concrete Postgres schema.Column into a sqlspec.Column.
-func columnSpec(col *schema.Column) (*sqlspec.Column, error) {
+func columnSpec(col *schema.Column, t *schema.Table) (*sqlspec.Column, error) {
 	ct, err := columnTypeSpec(col.Type.Type)
 	if err != nil {
 		return nil, err
diff --git a/sql/sqlite/sqlspec.go b/sql/sqlite/sqlspec.go
index 1e7b3ad7..09e7d817 100644
--- a/sql/sqlite/sqlspec.go
+++ b/sql/sqlite/sqlspec.go
@@ -193,7 +193,7 @@ func tableSpec(tab *schema.Table) (*sqlspec.Table, error) {
 }
 
 // columnSpec converts from a concrete SQLite schema.Column into a sqlspec.Column.
-func columnSpec(col *schema.Column) (*sqlspec.Column, error) {
+func columnSpec(col *schema.Column, t *schema.Table) (*sqlspec.Column, error) {
 	ct, err := columnTypeSpec(col.Type.Type)
 	if err != nil {
 		return nil, err


Test patch:
diff --git a/internal/integration/integration_test.go b/internal/integration/integration_test.go
index 14bd9b95..c55930a1 100644
--- a/internal/integration/integration_test.go
+++ b/internal/integration/integration_test.go
@@ -115,14 +115,14 @@ func testHCLIntegration(t T, full string, empty string) {
 
 func testCLISchemaInspect(t T, h string, dsn string, unmarshalSpec func(data []byte, unmarshaler schemaspec.Unmarshaler, v interface{}) error) {
 	// Required to have a clean "stderr" while running first time.
-	c := exec.Command("go", "run", "-mod=mod", "ariga.io/atlas/cmd/atlas")
-	require.NoError(t, c.Run())
+	err := exec.Command("go", "run", "-mod=mod", "ariga.io/atlas/cmd/atlas").Run()
+	require.NoError(t, err)
 	t.dropTables("users")
 	var expected schema.Schema
-	err := unmarshalSpec([]byte(h), schemahcl.Unmarshal, &expected)
+	err = unmarshalSpec([]byte(h), schemahcl.Unmarshal, &expected)
 	require.NoError(t, err)
 	t.applyHcl(h)
-	cmd := exec.Command("go", "run", "-mod=mod", "ariga.io/atlas/cmd/atlas",
+	cmd := exec.Command("go", "run", "ariga.io/atlas/cmd/atlas",
 		"schema",
 		"inspect",
 		"-d",
@@ -141,14 +141,14 @@ func testCLISchemaInspect(t T, h string, dsn string, unmarshalSpec func(data []b
 
 func testCLISchemaApply(t T, h string, dsn string) {
 	// Required to have a clean "stderr" while running first time.
-	c := exec.Command("go", "run", "-mod=mod", "ariga.io/atlas/cmd/atlas")
-	require.NoError(t, c.Run())
+	err := exec.Command("go", "run", "-mod=mod", "ariga.io/atlas/cmd/atlas").Run()
+	require.NoError(t, err)
 	t.dropTables("users")
 	f := "atlas.hcl"
-	err := ioutil.WriteFile(f, []byte(h), 0644)
+	err = ioutil.WriteFile(f, []byte(h), 0644)
 	require.NoError(t, err)
 	defer os.Remove(f)
-	cmd := exec.Command("go", "run", "-mod=mod", "ariga.io/atlas/cmd/atlas",
+	cmd := exec.Command("go", "run", "ariga.io/atlas/cmd/atlas",
 		"schema",
 		"apply",
 		"-d",
@@ -167,7 +167,7 @@ func testCLISchemaApply(t T, h string, dsn string) {
 	require.NoError(t, cmd.Run(), stderr.String(), stdout.String())
 	require.Empty(t, stderr.String(), stderr.String())
 	require.Contains(t, stdout.String(), "-- Planned")
-	u := t.users()
+	u := t.loadUsers()
 	require.NotNil(t, u)
 }
 
diff --git a/internal/integration/mysql_test.go b/internal/integration/mysql_test.go
index acb6adcb..54df771c 100644
--- a/internal/integration/mysql_test.go
+++ b/internal/integration/mysql_test.go
@@ -497,6 +497,8 @@ schema "test" {
 func TestMySQL_CLI(t *testing.T) {
 	h := `
 			schema "test" {
+				charset   = "%s"
+				collation = "%s"
 			}
 			table "users" {
 				schema = schema.test
@@ -509,12 +511,16 @@ func TestMySQL_CLI(t *testing.T) {
 			}`
 	t.Run("SchemaInspect", func(t *testing.T) {
 		myRun(t, func(t *myTest) {
-			testCLISchemaInspect(t, h, t.dsn(), mysql.UnmarshalSpec)
+			attrs := t.defaultAttrs()
+			charset, collate := attrs[0].(*schema.Charset), attrs[1].(*schema.Collation)
+			testCLISchemaInspect(t, fmt.Sprintf(h, charset.V, collate.V), t.dsn(), mysql.UnmarshalSpec)
 		})
 	})
 	t.Run("SchemaApply", func(t *testing.T) {
 		myRun(t, func(t *myTest) {
-			testCLISchemaApply(t, h, t.dsn())
+			attrs := t.defaultAttrs()
+			charset, collate := attrs[0].(*schema.Charset), attrs[1].(*schema.Collation)
+			testCLISchemaApply(t, fmt.Sprintf(h, charset.V, collate.V), t.dsn())
 		})
 	})
 }
diff --git a/sql/mysql/sqlspec_test.go b/sql/mysql/sqlspec_test.go
index a33eecb8..62a75234 100644
--- a/sql/mysql/sqlspec_test.go
+++ b/sql/mysql/sqlspec_test.go
@@ -60,7 +60,7 @@ table "accounts" {
 	}
 }
 `
-	s := schema.Schema{}
+	var s schema.Schema
 	err := UnmarshalSpec([]byte(f), schemahcl.Unmarshal, &s)
 	require.NoError(t, err)
 
@@ -151,6 +151,134 @@ table "accounts" {
 	require.EqualValues(t, exp, &s)
 }
 
+func TestMarshalSpec_Charset(t *testing.T) {
+	s := &schema.Schema{
+		Name: "test",
+		Attrs: []schema.Attr{
+			&schema.Charset{V: "utf8mb4"},
+			&schema.Collation{V: "utf8mb4_0900_ai_ci"},
+		},
+		Tables: []*schema.Table{
+			{
+				Name: "users",
+				Attrs: []schema.Attr{
+					&schema.Charset{V: "utf8mb4"},
+					&schema.Collation{V: "utf8mb4_0900_ai_ci"},
+				},
+				Columns: []*schema.Column{
+					{
+						Name: "a",
+						Type: &schema.ColumnType{Type: &schema.StringType{T: "text"}},
+						Attrs: []schema.Attr{
+							&schema.Charset{V: "latin1"},
+							&schema.Collation{V: "latin1_swedish_ci"},
+						},
+					},
+					{
+						Name: "b",
+						Type: &schema.ColumnType{Type: &schema.StringType{T: "text"}},
+						Attrs: []schema.Attr{
+							&schema.Charset{V: "utf8mb4"},
+							&schema.Collation{V: "utf8mb4_0900_ai_ci"},
+						},
+					},
+				},
+			},
+			{
+				Name: "posts",
+				Attrs: []schema.Attr{
+					&schema.Charset{V: "latin1"},
+					&schema.Collation{V: "latin1_swedish_ci"},
+				},
+				Columns: []*schema.Column{
+					{
+						Name: "a",
+						Type: &schema.ColumnType{Type: &schema.StringType{T: "text"}},
+						Attrs: []schema.Attr{
+							&schema.Charset{V: "latin1"},
+							&schema.Collation{V: "latin1_swedish_ci"},
+						},
+					},
+					{
+						Name: "b",
+						Type: &schema.ColumnType{Type: &schema.StringType{T: "text"}},
+						Attrs: []schema.Attr{
+							&schema.Charset{V: "utf8mb4"},
+							&schema.Collation{V: "utf8mb4_0900_ai_ci"},
+						},
+					},
+				},
+			},
+		},
+	}
+	s.Tables[0].Schema = s
+	s.Tables[1].Schema = s
+	buf, err := MarshalSpec(s, schemahcl.Marshal)
+	require.NoError(t, err)
+	// Charset and collate that are identical to their parent elements
+	// should not be printed as they are inherited by default from it.
+	const expected = `table "users" {
+  schema = schema.test
+  column "a" {
+    null      = false
+    type      = "text"
+    charset   = "latin1"
+    collation = "latin1_swedish_ci"
+  }
+  column "b" {
+    null = false
+    type = "text"
+  }
+}
+table "posts" {
+  schema    = schema.test
+  charset   = "latin1"
+  collation = "latin1_swedish_ci"
+  column "a" {
+    null = false
+    type = "text"
+  }
+  column "b" {
+    null      = false
+    type      = "text"
+    charset   = "utf8mb4"
+    collation = "utf8mb4_0900_ai_ci"
+  }
+}
+schema "test" {
+  charset   = "utf8mb4"
+  collation = "utf8mb4_0900_ai_ci"
+}
+`
+	require.EqualValues(t, expected, buf)
+
+	var (
+		s2    schema.Schema
+		latin = []schema.Attr{
+			&schema.Charset{V: "latin1"},
+			&schema.Collation{V: "latin1_swedish_ci"},
+		}
+		utf8mb4 = []schema.Attr{
+			&schema.Charset{V: "utf8mb4"},
+			&schema.Collation{V: "utf8mb4_0900_ai_ci"},
+		}
+	)
+	require.NoError(t, UnmarshalSpec(buf, schemahcl.Unmarshal, &s2))
+	require.Equal(t, utf8mb4, s2.Attrs)
+	posts, ok := s2.Table("posts")
+	require.True(t, ok)
+	require.Equal(t, latin, posts.Attrs)
+	users, ok := s2.Table("users")
+	require.True(t, ok)
+	require.Empty(t, users.Attrs)
+	a, ok := users.Column("a")
+	require.True(t, ok)
+	require.Equal(t, latin, a.Attrs)
+	b, ok := posts.Column("b")
+	require.True(t, ok)
+	require.Equal(t, utf8mb4, b.Attrs)
+}
+
 func TestUnmarshalSpecColumnTypes(t *testing.T) {
 	for _, tt := range []struct {
 		spec     *sqlspec.Column