Repository: ariga/atlas
Pull request #703 (2022-04-13 14:13:53)

sql/mysql: support unmarshaling generated columns


Patch:
diff --git a/sql/mysql/sqlspec.go b/sql/mysql/sqlspec.go
index 5ed604ae..2710b4cc 100644
--- a/sql/mysql/sqlspec.go
+++ b/sql/mysql/sqlspec.go
@@ -68,6 +68,7 @@ var (
 	hclState = schemahcl.New(
 		schemahcl.WithTypes(TypeRegistry.Specs()),
 		schemahcl.WithScopedEnums("table.index.type", IndexTypeBTree, IndexTypeHash, IndexTypeFullText, IndexTypeSpatial),
+		schemahcl.WithScopedEnums("table.column.as.type", stored, persistent, virtual),
 		schemahcl.WithScopedEnums("table.foreign_key.on_update", specutil.ReferenceVars...),
 		schemahcl.WithScopedEnums("table.foreign_key.on_delete", specutil.ReferenceVars...),
 	)
@@ -175,6 +176,9 @@ func convertColumn(spec *sqlspec.Column, _ *schema.Table) (*schema.Column, error
 			c.AddAttrs(&AutoIncrement{})
 		}
 	}
+	if err := convertGenExpr(spec.Remain(), c); err != nil {
+		return nil, err
+	}
 	return c, err
 }
 
@@ -277,6 +281,40 @@ func fromGenExpr(x schema.GeneratedExpr) *schemaspec.Resource {
 	}
 }
 
+func convertGenExpr(r *schemaspec.Resource, c *schema.Column) error {
+	asA, okA := r.Attr("as")
+	asR, okR := r.Resource("as")
+	switch {
+	case okA && okR:
+		return fmt.Errorf("multiple as definitions for column %q", c.Name)
+	case okA:
+		expr, err := asA.String()
+		if err != nil {
+			return err
+		}
+		c.Attrs = append(c.Attrs, &schema.GeneratedExpr{
+			Type: virtual,
+			Expr: expr,
+		})
+	case okR:
+		var spec struct {
+			Expr string `spec:"expr"`
+			Type string `spec:"type"`
+		}
+		if err := asR.As(&spec); err != nil {
+			return err
+		}
+		if spec.Type == "" {
+			spec.Type = virtual
+		}
+		c.Attrs = append(c.Attrs, &schema.GeneratedExpr{
+			Type: spec.Type,
+			Expr: spec.Expr,
+		})
+	}
+	return nil
+}
+
 // checkSpec converts from a concrete MySQL schema.Check into a sqlspec.Check.
 func checkSpec(s *schema.Check) *sqlspec.Check {
 	c := specutil.FromCheck(s)


Test patch:
diff --git a/sql/mysql/sqlspec_test.go b/sql/mysql/sqlspec_test.go
index b0bf2047..4886b208 100644
--- a/sql/mysql/sqlspec_test.go
+++ b/sql/mysql/sqlspec_test.go
@@ -723,6 +723,51 @@ schema "test" {
 	require.EqualValues(t, expected, string(buf))
 }
 
+func TestUnmarshalSpec_GeneratedColumns(t *testing.T) {
+	var (
+		s schema.Schema
+		f = `
+schema "test" {}
+table "users" {
+	schema = schema.test
+	column "c1" {
+		type = int
+	}
+	column "c2" {
+		type = int
+		as = "c1 * 2"
+	}
+	column "c3" {
+		type = int
+		as {
+			expr = "c2 * 2"
+		}
+	}
+	column "c4" {
+		type = int
+		as {
+			expr = "c3 * 2"
+			type = STORED
+		}
+	}
+}
+`
+	)
+	err := UnmarshalHCL([]byte(f), &s)
+	require.NoError(t, err)
+	exp := schema.New("test").
+		AddTables(
+			schema.NewTable("users").
+				AddColumns(
+					schema.NewIntColumn("c1", "int"),
+					schema.NewIntColumn("c2", "int").SetGeneratedExpr(&schema.GeneratedExpr{Expr: "c1 * 2", Type: "VIRTUAL"}),
+					schema.NewIntColumn("c3", "int").SetGeneratedExpr(&schema.GeneratedExpr{Expr: "c2 * 2", Type: "VIRTUAL"}),
+					schema.NewIntColumn("c4", "int").SetGeneratedExpr(&schema.GeneratedExpr{Expr: "c3 * 2", Type: "STORED"}),
+				),
+		)
+	require.EqualValues(t, exp, &s)
+}
+
 func TestMarshalSpec_FloatUnsigned(t *testing.T) {
 	s := schema.New("test").
 		AddTables(