Repository: ariga/atlas
Pull request #200 (2021-11-01 12:48:09)

sql/postgres: marshal spec


Patch:
diff --git a/sql/postgres/sqlspec.go b/sql/postgres/sqlspec.go
index 154f2a49..17f8c66b 100644
--- a/sql/postgres/sqlspec.go
+++ b/sql/postgres/sqlspec.go
@@ -3,6 +3,7 @@ package postgres
 import (
 	"errors"
 	"fmt"
+	"strconv"
 	"strings"
 
 	"ariga.io/atlas/schema/schemaspec"
@@ -37,6 +38,25 @@ func UnmarshalSpec(data []byte, unmarshaler schemaspec.Unmarshaler, v interface{
 	return nil
 }
 
+// MarshalSpec marshals v into an Atlas DDL document using a schemaspec.Marshaler.
+func MarshalSpec(v interface{}, marshaler schemaspec.Marshaler) ([]byte, error) {
+	var (
+		s  *schema.Schema
+		ok bool
+	)
+	if s, ok = v.(*schema.Schema); !ok {
+		return nil, fmt.Errorf("failed marshaling spec. %T is not supported", v)
+	}
+	spec, tables, err := schemaSpec(s)
+	if err != nil {
+		return nil, fmt.Errorf("failed converting schema to spec: %w", err)
+	}
+	return marshaler.MarshalSpec(&doc{
+		Tables:  tables,
+		Schemas: []*sqlspec.Schema{spec},
+	})
+}
+
 // convertTable converts a sqlspec.Table to a schema.Table. Table conversion is done without converting
 // 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.
@@ -188,3 +208,132 @@ func nparseRawType(spec *sqlspec.Column) (schema.Type, error) {
 	}
 	return columnType(cm), nil
 }
+
+// schemaSpec converts from a concrete Postgres schema to Atlas specification.
+func schemaSpec(schem *schema.Schema) (*sqlspec.Schema, []*sqlspec.Table, error) {
+	return specutil.FromSchema(schem, tableSpec)
+}
+
+// tableSpec converts from a concrete Postgres 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)
+}
+
+// columnSpec converts from a concrete Postgres schema.Column into a sqlspec.Column.
+func columnSpec(col *schema.Column) (*sqlspec.Column, error) {
+	ct, err := ncolumnTypeSpec(col.Type.Type)
+	if err != nil {
+		return nil, err
+	}
+	return &sqlspec.Column{
+		Name:     col.Name,
+		TypeName: ct.TypeName,
+		Null:     ct.Null,
+		DefaultExtension: schemaspec.DefaultExtension{
+			Extra: schemaspec.Resource{Attrs: ct.DefaultExtension.Extra.Attrs},
+		},
+	}, nil
+}
+
+// columnTypeSpec converts from a concrete Postgres schema.Type into sqlspec.Column Type.
+func ncolumnTypeSpec(t schema.Type) (*sqlspec.Column, error) {
+	switch t := t.(type) {
+	case *schema.EnumType:
+		return nenumSpec(t)
+	case *schema.IntegerType:
+		return nintegerSpec(t)
+	case *schema.StringType:
+		return nstringSpec(t)
+	case *schema.DecimalType:
+		precision := specutil.LitAttr("precision", strconv.Itoa(t.Precision))
+		scale := specutil.LitAttr("scale", strconv.Itoa(t.Scale))
+		return specutil.NewCol("", "decimal", precision, scale), nil
+	case *schema.BinaryType:
+		return &sqlspec.Column{TypeName: t.T}, nil
+	case *schema.BoolType:
+		return &sqlspec.Column{TypeName: "boolean"}, nil
+	case *schema.FloatType:
+		precision := specutil.LitAttr("precision", strconv.Itoa(t.Precision))
+		return specutil.NewCol("", "float", precision), nil
+	case *schema.TimeType:
+		return &sqlspec.Column{TypeName: t.T}, nil
+	case *schema.JSONType:
+		return &sqlspec.Column{TypeName: t.T}, nil
+	case *schema.SpatialType:
+		return &sqlspec.Column{TypeName: t.T}, nil
+	case *schema.UnsupportedType:
+		return &sqlspec.Column{TypeName: t.T}, nil
+	case *NetworkType:
+		return &sqlspec.Column{TypeName: t.T}, nil
+	case *CurrencyType:
+		return &sqlspec.Column{TypeName: t.T}, nil
+	case *BitType:
+		return nbitSpec(t)
+	default:
+		return nil, fmt.Errorf("failed to convert column type %T to spec", t)
+	}
+}
+
+// temporarily prefixed with "n" until we complete the refactor of replacing sql/schemaspec with sqlspec.
+func nstringSpec(t *schema.StringType) (*sqlspec.Column, error) {
+	switch t.T {
+	case tVarChar, tText:
+		s := strconv.Itoa(t.Size)
+		return specutil.NewCol("", "string", specutil.LitAttr("size", s)), nil
+	default:
+		return nil, errors.New("schema string failed to convert")
+	}
+}
+
+// temporarily prefixed with "n" until we complete the refactor of replacing sql/schemaspec with sqlspec.
+func nintegerSpec(t *schema.IntegerType) (*sqlspec.Column, error) {
+	switch t.T {
+	case tInt:
+		if t.Unsigned {
+			return specutil.NewCol("", "uint"), nil
+		}
+		return &sqlspec.Column{TypeName: "int"}, nil
+	case tBigInt:
+		if t.Unsigned {
+			return specutil.NewCol("", "uint64"), nil
+		}
+		return &sqlspec.Column{TypeName: "int64"}, nil
+	default:
+		return nil, errors.New("schema integer failed to convert")
+	}
+}
+
+// temporarily prefixed with "n" until we complete the refactor of replacing sql/schemaspec with sqlspec.
+func nenumSpec(t *schema.EnumType) (*sqlspec.Column, error) {
+	if len(t.Values) == 0 {
+		return nil, errors.New("schema enum fields to have values")
+	}
+	quoted := make([]string, 0, len(t.Values))
+	for _, v := range t.Values {
+		quoted = append(quoted, strconv.Quote(v))
+	}
+	return specutil.NewCol("", "enum", specutil.ListAttr("values", quoted...)), nil
+}
+
+// temporarily prefixed with "n" until we complete the refactor of replacing sql/schemaspec with sqlspec.
+func nbitSpec(t *BitType) (*sqlspec.Column, error) {
+	var c *sqlspec.Column
+	switch t.T {
+	case tBit:
+		if t.Len == 1 {
+			c = &sqlspec.Column{TypeName: tBit}
+		} else {
+			c = specutil.NewCol("", fmt.Sprintf("%s(%d)", tBit, t.Len))
+		}
+		return c, nil
+	case tBitVar:
+		if t.Len == 0 {
+			c = &sqlspec.Column{TypeName: tBitVar}
+		} else {
+			c = specutil.NewCol("", fmt.Sprintf("%s(%d)", tBitVar, t.Len))
+		}
+		return c, nil
+	default:
+		return nil, errors.New("schema bit failed to convert")
+	}
+}


Test patch:
diff --git a/sql/postgres/sqlspec_test.go b/sql/postgres/sqlspec_test.go
index 65140254..1948bd11 100644
--- a/sql/postgres/sqlspec_test.go
+++ b/sql/postgres/sqlspec_test.go
@@ -157,35 +157,35 @@ func TestUnmarshalSpecColumnTypes(t *testing.T) {
 		{
 			spec: specutil.NewCol("int64", "int64"),
 			expected: &schema.IntegerType{
-				T:        "bigint",
+				T:        tBigInt,
 				Unsigned: false,
 			},
 		},
 		{
 			spec: specutil.NewCol("string_varchar", "string", specutil.LitAttr("size", "255")),
 			expected: &schema.StringType{
-				T:    "varchar",
+				T:    tVarChar,
 				Size: 255,
 			},
 		},
 		{
 			spec: specutil.NewCol("string_test", "string", specutil.LitAttr("size", "10485761")),
 			expected: &schema.StringType{
-				T:    "text",
+				T:    tText,
 				Size: 10_485_761,
 			},
 		},
 		{
 			spec: specutil.NewCol("varchar(255)", "varchar(255)"),
 			expected: &schema.StringType{
-				T:    "varchar",
+				T:    tVarChar,
 				Size: 255,
 			},
 		},
 		{
 			spec: specutil.NewCol("decimal(10, 2)", "decimal(10, 2)"),
 			expected: &schema.DecimalType{
-				T:         "decimal",
+				T:         tDecimal,
 				Scale:     2,
 				Precision: 10,
 			},
@@ -196,43 +196,43 @@ func TestUnmarshalSpecColumnTypes(t *testing.T) {
 		},
 		{
 			spec:     specutil.NewCol("bool", "boolean"),
-			expected: &schema.BoolType{T: "boolean"},
+			expected: &schema.BoolType{T: tBoolean},
 		},
 		{
 			spec:     specutil.NewCol("decimal", "decimal", specutil.LitAttr("precision", "10"), specutil.LitAttr("scale", "2")),
-			expected: &schema.DecimalType{T: "decimal", Precision: 10, Scale: 2},
+			expected: &schema.DecimalType{T: tDecimal, Precision: 10, Scale: 2},
 		},
 		{
 			spec:     specutil.NewCol("float", "float", specutil.LitAttr("precision", "10")),
-			expected: &schema.FloatType{T: "real", Precision: 10},
+			expected: &schema.FloatType{T: tReal, Precision: 10},
 		},
 		{
 			spec:     specutil.NewCol("float", "float", specutil.LitAttr("precision", "25")),
-			expected: &schema.FloatType{T: "double precision", Precision: 25},
+			expected: &schema.FloatType{T: tDouble, Precision: 25},
 		},
 		{
 			spec:     specutil.NewCol("cidr", "cidr"),
-			expected: &NetworkType{T: "cidr"},
+			expected: &NetworkType{T: tCIDR},
 		},
 		{
 			spec:     specutil.NewCol("money", "money"),
-			expected: &CurrencyType{T: "money"},
+			expected: &CurrencyType{T: tMoney},
 		},
 		{
 			spec:     specutil.NewCol("bit", "bit"),
-			expected: &BitType{T: "bit", Len: 1},
+			expected: &BitType{T: tBit, Len: 1},
 		},
 		{
 			spec:     specutil.NewCol("bitvar", "bit varying"),
-			expected: &BitType{T: "bit varying"},
+			expected: &BitType{T: tBitVar},
 		},
 		{
 			spec:     specutil.NewCol("bitvar8", "bit varying(8)"),
-			expected: &BitType{T: "bit varying", Len: 8},
+			expected: &BitType{T: tBitVar, Len: 8},
 		},
 		{
 			spec:     specutil.NewCol("bit8", "bit(8)"),
-			expected: &BitType{T: "bit", Len: 8},
+			expected: &BitType{T: tBit, Len: 8},
 		},
 	} {
 		t.Run(tt.spec.Name, func(t *testing.T) {
@@ -292,3 +292,140 @@ table "table" {
 	body := fmt.Sprintf(tmpl, buf)
 	return []byte(body)
 }
+
+func TestMarshalSpecColumnType(t *testing.T) {
+	for _, tt := range []struct {
+		schem    schema.Type
+		expected *sqlspec.Column
+	}{
+		{
+			schem: &schema.IntegerType{
+				T:        tInt,
+				Unsigned: false,
+			},
+			expected: specutil.NewCol("int", "int"),
+		},
+		{
+			schem: &schema.IntegerType{
+				T:        tBigInt,
+				Unsigned: false,
+			},
+			expected: specutil.NewCol("int64", "int64"),
+		},
+		{
+			schem: &schema.StringType{
+				T:    tVarChar,
+				Size: 255,
+			},
+			expected: specutil.NewCol("string_varchar", "string", specutil.LitAttr("size", "255")),
+		},
+		{
+			schem: &schema.StringType{
+				T:    tText,
+				Size: 10_485_761,
+			},
+			expected: specutil.NewCol("string_text", "string", specutil.LitAttr("size", "10485761")),
+		},
+		{
+			schem: &schema.DecimalType{
+				T:         tDecimal,
+				Scale:     2,
+				Precision: 10,
+			},
+			expected: specutil.NewCol("decimal", "decimal",
+				specutil.LitAttr("precision", "10"), specutil.LitAttr("scale", "2")),
+		},
+		{
+			schem: &schema.EnumType{
+				Values: []string{"a", "b", "c"},
+			},
+			expected: specutil.NewCol("enum", "enum",
+				specutil.ListAttr("values", "a", "b", "c")),
+		},
+		{
+			schem: &schema.BoolType{
+				T: tBoolean,
+			},
+			expected: specutil.NewCol("boolean", "boolean"),
+		},
+		{
+			schem: &schema.FloatType{
+				T:         tReal,
+				Precision: 10,
+			},
+			expected: specutil.NewCol("float_real", "float", specutil.LitAttr("precision", "10")),
+		},
+		{
+			schem: &schema.FloatType{
+				T:         tDouble,
+				Precision: 25,
+			},
+			expected: specutil.NewCol("float_double", "float", specutil.LitAttr("precision", "25")),
+		},
+		{
+			schem: &NetworkType{
+				T: tCIDR,
+			},
+			expected: specutil.NewCol("network", "cidr"),
+		},
+		{
+			schem: &CurrencyType{
+				T: tMoney,
+			},
+			expected: specutil.NewCol("money", "money"),
+		},
+		{
+			schem: &BitType{
+				T: tBit,
+				Len: 1,
+			},
+			expected: specutil.NewCol("bit", "bit"),
+		},
+		{
+			schem: &BitType{
+				T: tBitVar,
+			},
+			expected: specutil.NewCol("bitvar", "bit varying"),
+		},
+		{
+			schem: &BitType{
+				T: tBit,
+				Len: 8,
+			},
+			expected: specutil.NewCol("bit8", "bit(8)"),
+		},
+		{
+			schem: &BitType{
+				T: tBitVar,
+				Len: 8,
+			},
+			expected: specutil.NewCol("bitvar8", "bit varying(8)"),
+		},
+	} {
+		t.Run(tt.expected.Name, func(t *testing.T) {
+			s := schema.Schema{
+				Tables: []*schema.Table{
+					{
+						Name: "table",
+						Columns: []*schema.Column{
+							{
+								Name: "column",
+								Type: &schema.ColumnType{Type: tt.schem},
+							},
+						},
+					},
+				},
+			}
+			s.Tables[0].Schema = &s
+			ddl, err := MarshalSpec(&s, schemahcl.Marshal)
+			require.NoError(t, err)
+			var test struct {
+				Table *sqlspec.Table `spec:"table"`
+			}
+			err = schemahcl.Unmarshal(ddl, &test)
+			require.NoError(t, err)
+			require.EqualValues(t, tt.expected.TypeName, test.Table.Columns[0].TypeName)
+			require.ElementsMatch(t, tt.expected.Extra.Attrs, test.Table.Columns[0].Extra.Attrs)
+		})
+	}
+}