Repository: ariga/atlas
Pull request #723 (2022-04-18 20:29:19)

sql/postgres: support (un)marshal generated columns


Patch:
diff --git a/sql/internal/specutil/convert.go b/sql/internal/specutil/convert.go
index 5df942db..195f4b48 100644
--- a/sql/internal/specutil/convert.go
+++ b/sql/internal/specutil/convert.go
@@ -280,23 +280,6 @@ func linkForeignKeys(tbl *schema.Table, sch *schema.Schema, table *sqlspec.Table
 	return nil
 }
 
-// FromRealm converts a schema.Realm into []sqlspec.Schema and []sqlspec.Table.
-func FromRealm(r *schema.Realm, fn TableSpecFunc) ([]*sqlspec.Schema, []*sqlspec.Table, error) {
-	var (
-		schemas []*sqlspec.Schema
-		tables  []*sqlspec.Table
-	)
-	for _, sch := range r.Schemas {
-		schemaSpec, tableSpecs, err := FromSchema(sch, fn)
-		if err != nil {
-			return nil, nil, fmt.Errorf("specutil: cannot convert schema %q: %w", sch.Name, err)
-		}
-		schemas = append(schemas, schemaSpec)
-		tables = append(tables, tableSpecs...)
-	}
-	return schemas, tables, nil
-}
-
 // FromSchema converts a schema.Schema into sqlspec.Schema and []sqlspec.Table.
 func FromSchema(s *schema.Schema, fn TableSpecFunc) (*sqlspec.Schema, []*sqlspec.Table, error) {
 	spec := &sqlspec.Schema{
@@ -395,6 +378,49 @@ func FromColumn(col *schema.Column, columnTypeSpec ColumnTypeSpecFunc) (*sqlspec
 	return spec, nil
 }
 
+// FromGenExpr returns the spec for a generated expression.
+func FromGenExpr(x schema.GeneratedExpr, t func(string) string) *schemaspec.Resource {
+	return &schemaspec.Resource{
+		Type: "as",
+		Attrs: []*schemaspec.Attr{
+			StrAttr("expr", x.Expr),
+			VarAttr("type", t(x.Type)),
+		},
+	}
+}
+
+// ConvertGenExpr converts the "as" attribute or the block under the given resource.
+func ConvertGenExpr(r *schemaspec.Resource, c *schema.Column, t func(string) string) 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: t(""), // default type.
+			Expr: expr,
+		})
+	case okR:
+		var spec struct {
+			Expr string `spec:"expr"`
+			Type string `spec:"type"`
+		}
+		if err := asR.As(&spec); err != nil {
+			return err
+		}
+		c.Attrs = append(c.Attrs, &schema.GeneratedExpr{
+			Expr: spec.Expr,
+			Type: t(spec.Type),
+		})
+	}
+	return nil
+}
+
 func toValue(expr schema.Expr) (schemaspec.Value, error) {
 	var (
 		v   string
diff --git a/sql/mysql/sqlspec.go b/sql/mysql/sqlspec.go
index da0e96d2..5471bb3d 100644
--- a/sql/mysql/sqlspec.go
+++ b/sql/mysql/sqlspec.go
@@ -176,7 +176,7 @@ func convertColumn(spec *sqlspec.Column, _ *schema.Table) (*schema.Column, error
 			c.AddAttrs(&AutoIncrement{})
 		}
 	}
-	if err := convertGenExpr(spec.Remain(), c); err != nil {
+	if err := specutil.ConvertGenExpr(spec.Remain(), c, storedOrVirtual); err != nil {
 		return nil, err
 	}
 	return c, err
@@ -261,53 +261,11 @@ func columnSpec(c *schema.Column, t *schema.Table) (*sqlspec.Column, error) {
 		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))
+		spec.Extra.Children = append(spec.Extra.Children, specutil.FromGenExpr(x, storedOrVirtual))
 	}
 	return spec, nil
 }
 
-// fromGenExpr returns the spec for a generated expression.
-func fromGenExpr(x schema.GeneratedExpr) *schemaspec.Resource {
-	return &schemaspec.Resource{
-		Type: "as",
-		Attrs: []*schemaspec.Attr{
-			specutil.StrAttr("expr", x.Expr),
-			specutil.VarAttr("type", storedOrVirtual(x.Type)),
-		},
-	}
-}
-
-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
-		}
-		c.Attrs = append(c.Attrs, &schema.GeneratedExpr{
-			Expr: spec.Expr,
-			Type: storedOrVirtual(spec.Type),
-		})
-	}
-	return nil
-}
-
 // storedOrVirtual returns a STORED or VIRTUAL
 // generated type option based on the given string.
 func storedOrVirtual(s string) string {
diff --git a/sql/postgres/sqlspec.go b/sql/postgres/sqlspec.go
index d6e87200..8d5fb52b 100644
--- a/sql/postgres/sqlspec.go
+++ b/sql/postgres/sqlspec.go
@@ -107,6 +107,7 @@ var (
 		schemahcl.WithTypes(TypeRegistry.Specs()),
 		schemahcl.WithScopedEnums("table.index.type", IndexTypeBTree, IndexTypeHash, IndexTypeGIN, IndexTypeGiST),
 		schemahcl.WithScopedEnums("table.column.identity.generated", GeneratedTypeAlways, GeneratedTypeByDefault),
+		schemahcl.WithScopedEnums("table.column.as.type", "STORED"),
 		schemahcl.WithScopedEnums("table.foreign_key.on_update", specutil.ReferenceVars...),
 		schemahcl.WithScopedEnums("table.foreign_key.on_delete", specutil.ReferenceVars...),
 	)
@@ -143,6 +144,9 @@ func convertColumn(spec *sqlspec.Column, _ *schema.Table) (*schema.Column, error
 		}
 		c.Attrs = append(c.Attrs, id)
 	}
+	if err := specutil.ConvertGenExpr(spec.Remain(), c, generatedType); err != nil {
+		return nil, err
+	}
 	return c, nil
 }
 
@@ -345,6 +349,9 @@ func columnSpec(c *schema.Column, _ *schema.Table) (*sqlspec.Column, error) {
 	if i := (&Identity{}); sqlx.Has(c.Attrs, i) {
 		s.Extra.Children = append(s.Extra.Children, fromIdentity(i))
 	}
+	if x := (schema.GeneratedExpr{}); sqlx.Has(c.Attrs, &x) {
+		s.Extra.Children = append(s.Extra.Children, specutil.FromGenExpr(x, generatedType))
+	}
 	return s, nil
 }
 
@@ -488,3 +495,6 @@ func formatTime() schemahcl.TypeSpecOption {
 		return FormatType(&schema.TimeType{T: t.T, Precision: &p})
 	})
 }
+
+// generatedType returns the default and only type for a generated column.
+func generatedType(string) string { return "STORED" }


Test patch:
diff --git a/sql/postgres/sqlspec_test.go b/sql/postgres/sqlspec_test.go
index e84a2084..a85686ce 100644
--- a/sql/postgres/sqlspec_test.go
+++ b/sql/postgres/sqlspec_test.go
@@ -368,6 +368,85 @@ table "t" {
 	})
 }
 
+func TestMarshalSpec_GeneratedColumn(t *testing.T) {
+	s := schema.New("test").
+		AddTables(
+			schema.NewTable("users").
+				AddColumns(
+					schema.NewIntColumn("c1", "int").
+						SetGeneratedExpr(&schema.GeneratedExpr{Expr: "c1 * 2"}),
+					schema.NewIntColumn("c2", "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
+    as {
+      expr = "c1 * 2"
+      type = STORED
+    }
+  }
+  column "c2" {
+    null = false
+    type = int
+    as {
+      expr = "c3 * c4"
+      type = STORED
+    }
+  }
+}
+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
+		as = "1"
+	}
+	column "c2" {
+		type = int
+		as {
+			expr = "2"
+		}
+	}
+	column "c3" {
+		type = int
+		as {
+			expr = "3"
+			type = STORED
+		}
+	}
+}
+`
+	)
+	err := UnmarshalHCL([]byte(f), &s)
+	require.NoError(t, err)
+	exp := schema.New("test").
+		AddTables(
+			schema.NewTable("users").
+				AddColumns(
+					schema.NewIntColumn("c1", "int").SetGeneratedExpr(&schema.GeneratedExpr{Expr: "1", Type: "STORED"}),
+					schema.NewIntColumn("c2", "int").SetGeneratedExpr(&schema.GeneratedExpr{Expr: "2", Type: "STORED"}),
+					schema.NewIntColumn("c3", "int").SetGeneratedExpr(&schema.GeneratedExpr{Expr: "3", Type: "STORED"}),
+				),
+		)
+	require.EqualValues(t, exp, &s)
+}
+
 func TestMarshalSpec_Enum(t *testing.T) {
 	s := schema.New("test").
 		AddTables(