Repository: ariga/atlas
Pull request #701 (2022-04-13 10:02:10)

sql/mysql: support marshaling generated columns


Patch:
diff --git a/schema/schemaspec/schemahcl/hcl.go b/schema/schemaspec/schemahcl/hcl.go
index 0c53dc5e..71033f47 100644
--- a/schema/schemaspec/schemahcl/hcl.go
+++ b/schema/schemaspec/schemahcl/hcl.go
@@ -330,7 +330,7 @@ func (s *State) writeAttr(attr *schemaspec.Attr, body *hclwrite.Body) error {
 	return nil
 }
 
-// normalizeLiterals transforms attriburtes with LiteralValue that cannot be
+// normalizeLiterals transforms attributes with LiteralValue that cannot be
 // written as correct HCL into RawExpr.
 func normalizeLiterals(attr *schemaspec.Attr) *schemaspec.Attr {
 	lv, ok := attr.V.(*schemaspec.LiteralValue)
diff --git a/sql/mysql/driver.go b/sql/mysql/driver.go
index 437b60cf..80175271 100644
--- a/sql/mysql/driver.go
+++ b/sql/mysql/driver.go
@@ -282,7 +282,7 @@ const (
 	defaultGen    = "default_generated"
 	autoIncrement = "auto_increment"
 
-	genVirtual    = "virtual"
-	genStored     = "stored"
-	genPersistent = "persistent"
+	virtual    = "VIRTUAL"
+	stored     = "STORED"
+	persistent = "PERSISTENT"
 )
diff --git a/sql/mysql/sqlspec.go b/sql/mysql/sqlspec.go
index 0a4fe6c9..5ed604ae 100644
--- a/sql/mysql/sqlspec.go
+++ b/sql/mysql/sqlspec.go
@@ -240,23 +240,41 @@ func partAttr(part *schema.IndexPart, spec *sqlspec.IndexPart) {
 
 // columnSpec converts from a concrete MySQL schema.Column into a sqlspec.Column.
 func columnSpec(c *schema.Column, t *schema.Table) (*sqlspec.Column, error) {
-	col, err := specutil.FromColumn(c, columnTypeSpec)
+	spec, err := specutil.FromColumn(c, columnTypeSpec)
 	if err != nil {
 		return nil, err
 	}
 	if c, ok := hasCharset(c.Attrs, t.Attrs); ok {
-		col.Extra.Attrs = append(col.Extra.Attrs, specutil.StrAttr("charset", c))
+		spec.Extra.Attrs = append(spec.Extra.Attrs, specutil.StrAttr("charset", c))
 	}
 	if c, ok := hasCollate(c.Attrs, t.Attrs); ok {
-		col.Extra.Attrs = append(col.Extra.Attrs, specutil.StrAttr("collate", c))
+		spec.Extra.Attrs = append(spec.Extra.Attrs, specutil.StrAttr("collate", c))
 	}
 	if o := (OnUpdate{}); sqlx.Has(c.Attrs, &o) {
-		col.Extra.Attrs = append(col.Extra.Attrs, specutil.RawAttr("on_update", o.A))
+		spec.Extra.Attrs = append(spec.Extra.Attrs, specutil.RawAttr("on_update", o.A))
 	}
 	if sqlx.Has(c.Attrs, &AutoIncrement{}) {
-		col.Extra.Attrs = append(col.Extra.Attrs, specutil.BoolAttr("auto_increment", true))
+		spec.Extra.Attrs = append(spec.Extra.Attrs, specutil.BoolAttr("auto_increment", true))
+	}
+	if x := (schema.GeneratedExpr{}); sqlx.Has(c.Attrs, &x) {
+		spec.Extra.Children = append(spec.Extra.Children, fromGenExpr(x))
+	}
+	return spec, nil
+}
+
+// fromGenExpr returns the spec for a generated expression.
+func fromGenExpr(x schema.GeneratedExpr) *schemaspec.Resource {
+	t := strings.ToUpper(x.Type)
+	if t == "" {
+		t = virtual
+	}
+	return &schemaspec.Resource{
+		Type: "as",
+		Attrs: []*schemaspec.Attr{
+			specutil.StrAttr("expr", x.Expr),
+			specutil.VarAttr("type", t),
+		},
 	}
-	return col, nil
 }
 
 // checkSpec converts from a concrete MySQL schema.Check into a sqlspec.Check.
diff --git a/sql/sqlspec/sqlspec.go b/sql/sqlspec/sqlspec.go
index 0a053efa..8741a5e1 100644
--- a/sql/sqlspec/sqlspec.go
+++ b/sql/sqlspec/sqlspec.go
@@ -5,7 +5,6 @@ import (
 )
 
 type (
-
 	// Schema holds a specification for a Schema.
 	Schema struct {
 		Name string `spec:",name"`


Test patch:
diff --git a/sql/mysql/sqlspec_test.go b/sql/mysql/sqlspec_test.go
index 10ce039a..b0bf2047 100644
--- a/sql/mysql/sqlspec_test.go
+++ b/sql/mysql/sqlspec_test.go
@@ -670,6 +670,59 @@ schema "test" {
 	require.EqualValues(t, expected, string(buf))
 }
 
+func TestMarshalSpec_GeneratedColumn(t *testing.T) {
+	s := schema.New("test").
+		AddTables(
+			schema.NewTable("users").
+				AddColumns(
+					schema.NewIntColumn("c1", "int"),
+					schema.NewIntColumn("c2", "int").
+						SetGeneratedExpr(&schema.GeneratedExpr{Expr: "c1 * 2"}),
+					schema.NewIntColumn("c3", "int").
+						SetGeneratedExpr(&schema.GeneratedExpr{Expr: "c2 * c3", Type: "VIRTUAL"}),
+					schema.NewIntColumn("c4", "int").
+						SetGeneratedExpr(&schema.GeneratedExpr{Expr: "c3 * c4", Type: "STORED"}),
+				),
+		)
+	buf, err := MarshalSpec(s, hclState)
+	require.NoError(t, err)
+	const expected = `table "users" {
+  schema = schema.test
+  column "c1" {
+    null = false
+    type = int
+  }
+  column "c2" {
+    null = false
+    type = int
+    as {
+      expr = "c1 * 2"
+      type = VIRTUAL
+    }
+  }
+  column "c3" {
+    null = false
+    type = int
+    as {
+      expr = "c2 * c3"
+      type = VIRTUAL
+    }
+  }
+  column "c4" {
+    null = false
+    type = int
+    as {
+      expr = "c3 * c4"
+      type = STORED
+    }
+  }
+}
+schema "test" {
+}
+`
+	require.EqualValues(t, expected, string(buf))
+}
+
 func TestMarshalSpec_FloatUnsigned(t *testing.T) {
 	s := schema.New("test").
 		AddTables(