Repository: ariga/atlas
Pull request #198 (2021-10-31 14:56:31)

sql/postgres: unmarshal spec column types
- convert column type


Patch:
diff --git a/sql/postgres/sqlspec.go b/sql/postgres/sqlspec.go
new file mode 100644
index 00000000..154f2a49
--- /dev/null
+++ b/sql/postgres/sqlspec.go
@@ -0,0 +1,190 @@
+package postgres
+
+import (
+	"errors"
+	"fmt"
+	"strings"
+
+	"ariga.io/atlas/schema/schemaspec"
+	"ariga.io/atlas/sql/internal/specutil"
+	"ariga.io/atlas/sql/schema"
+	"ariga.io/atlas/sql/sqlspec"
+)
+
+type doc struct {
+	Tables  []*sqlspec.Table  `spec:"table"`
+	Schemas []*sqlspec.Schema `spec:"schema"`
+}
+
+// UnmarshalSpec unmarshals an Atlas DDL document using an unmarshaler into v.
+func UnmarshalSpec(data []byte, unmarshaler schemaspec.Unmarshaler, v interface{}) error {
+	var d doc
+	if err := unmarshaler.UnmarshalSpec(data, &d); err != nil {
+		return err
+	}
+	s, ok := v.(*schema.Schema)
+	if !ok {
+		return fmt.Errorf("postgres: failed unmarshaling spec. %T is not supported", v)
+	}
+	if len(d.Schemas) != 1 {
+		return fmt.Errorf("postgres: 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("postgres: failed converting to *schema.Schema: %w", err)
+	}
+	*s = *conv
+	return nil
+}
+
+// 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.
+func convertTable(spec *sqlspec.Table, parent *schema.Schema) (*schema.Table, error) {
+	return specutil.Table(spec, parent, convertColumn, convertPrimaryKey, convertIndex)
+}
+
+// convertPrimaryKey converts a sqlspec.PrimaryKey to a schema.Index.
+func convertPrimaryKey(spec *sqlspec.PrimaryKey, parent *schema.Table) (*schema.Index, error) {
+	return specutil.PrimaryKey(spec, parent)
+}
+
+// convertIndex converts an sqlspec.Index to a schema.Index.
+func convertIndex(spec *sqlspec.Index, parent *schema.Table) (*schema.Index, error) {
+	return specutil.Index(spec, parent)
+}
+
+// convertColumn converts a sqlspec.Column into a schema.Column.
+func convertColumn(spec *sqlspec.Column, _ *schema.Table) (*schema.Column, error) {
+	return specutil.Column(spec, convertColumnType)
+}
+
+// convertColumnType converts a sqlspec.Column into a concrete Postgres schema.Type.
+func convertColumnType(spec *sqlspec.Column) (schema.Type, error) {
+	switch sqlspec.Type(spec.TypeName) {
+	case sqlspec.TypeInt, sqlspec.TypeInt8, sqlspec.TypeInt16,
+		sqlspec.TypeInt64, sqlspec.TypeUint, sqlspec.TypeUint8,
+		sqlspec.TypeUint16, sqlspec.TypeUint64:
+		return nconvertInteger(spec)
+	case sqlspec.TypeString:
+		return nconvertString(spec)
+	case sqlspec.TypeEnum:
+		return nconvertEnum(spec)
+	case sqlspec.TypeDecimal:
+		return nconvertDecimal(spec)
+	case sqlspec.TypeFloat:
+		return nconvertFloat(spec)
+	case sqlspec.TypeTime:
+		return &schema.TimeType{T: tTimestamp}, nil
+	case sqlspec.TypeBinary:
+		return &schema.BinaryType{T: tBytea}, nil
+	case sqlspec.TypeBoolean:
+		return &schema.BoolType{T: tBoolean}, nil
+	default:
+		return nparseRawType(spec)
+	}
+}
+
+// temporarily prefixed with "n" until we complete the refactor of replacing sql/schemaspec with sqlspec.
+func nconvertInteger(spec *sqlspec.Column) (schema.Type, error) {
+	if strings.HasPrefix(spec.TypeName, "u") {
+		return nil, fmt.Errorf("unsigned integers currently not supported")
+	}
+	typ := &schema.IntegerType{}
+	switch sqlspec.Type(spec.TypeName) {
+	case sqlspec.TypeInt8:
+		return nil, fmt.Errorf("8-bit integers not supported")
+	case sqlspec.TypeInt16:
+		typ.T = tSmallInt
+	case sqlspec.TypeInt:
+		typ.T = tInteger
+	case sqlspec.TypeInt64:
+		typ.T = tBigInt
+	default:
+		return nil, fmt.Errorf("unknown integer column type %q", spec.TypeName)
+	}
+	return typ, nil
+}
+
+// temporarily prefixed with "n" until we complete the refactor of replacing sql/schemaspec with sqlspec.
+func nconvertString(spec *sqlspec.Column) (schema.Type, error) {
+	st := &schema.StringType{
+		Size: 255,
+	}
+	if attr, ok := spec.Attr("size"); ok {
+		s, err := attr.Int()
+		if err != nil {
+			return nil, err
+		}
+		st.Size = s
+	}
+	switch {
+	case st.Size < maxCharSize:
+		st.T = tVarChar
+	default:
+		st.T = tText
+	}
+	return st, nil
+}
+
+// temporarily prefixed with "n" until we complete the refactor of replacing sql/schemaspec with sqlspec.
+func nconvertEnum(spec *sqlspec.Column) (schema.Type, error) {
+	attr, ok := spec.Attr("values")
+	if !ok {
+		return nil, errors.New("expected enum fields to have values")
+	}
+	list, err := attr.Strings()
+	if err != nil {
+		return nil, err
+	}
+	return &schema.EnumType{Values: list}, nil
+}
+
+// temporarily prefixed with "n" until we complete the refactor of replacing sql/schemaspec with sqlspec.
+func nconvertDecimal(spec *sqlspec.Column) (schema.Type, error) {
+	dt := &schema.DecimalType{
+		T: tDecimal,
+	}
+	if precision, ok := spec.Attr("precision"); ok {
+		p, err := precision.Int()
+		if err != nil {
+			return nil, err
+		}
+		dt.Precision = p
+	}
+	if scale, ok := spec.Attr("scale"); ok {
+		s, err := scale.Int()
+		if err != nil {
+			return nil, err
+		}
+		dt.Scale = s
+	}
+	return dt, nil
+}
+
+// temporarily prefixed with "n" until we complete the refactor of replacing sql/schemaspec with sqlspec.
+func nconvertFloat(spec *sqlspec.Column) (schema.Type, error) {
+	ft := &schema.FloatType{
+		T: tReal,
+	}
+	if precision, ok := spec.Attr("precision"); ok {
+		p, err := precision.Int()
+		if err != nil {
+			return nil, err
+		}
+		ft.Precision = p
+	}
+	if ft.Precision > 23 {
+		ft.T = tDouble
+	}
+	return ft, nil
+}
+
+// temporarily prefixed with "n" until we complete the refactor of replacing sql/schemaspec with sqlspec.
+func nparseRawType(spec *sqlspec.Column) (schema.Type, error) {
+	cm, err := parseColumn(spec.TypeName)
+	if err != nil {
+		return nil, err
+	}
+	return columnType(cm), nil
+}


Test patch:
diff --git a/sql/postgres/sqlspec_test.go b/sql/postgres/sqlspec_test.go
new file mode 100644
index 00000000..77b3da80
--- /dev/null
+++ b/sql/postgres/sqlspec_test.go
@@ -0,0 +1,158 @@
+package postgres
+
+import (
+	"fmt"
+	"log"
+	"testing"
+
+	"ariga.io/atlas/schema/schemaspec/schemahcl"
+	"ariga.io/atlas/sql/internal/specutil"
+	"ariga.io/atlas/sql/schema"
+	"ariga.io/atlas/sql/sqlspec"
+
+	"github.com/stretchr/testify/require"
+)
+
+func TestUnmarshalSpecColumnTypes(t *testing.T) {
+	for _, tt := range []struct {
+		spec     *sqlspec.Column
+		expected schema.Type
+	}{
+		{
+			spec: specutil.NewCol("int64", "int64"),
+			expected: &schema.IntegerType{
+				T:        "bigint",
+				Unsigned: false,
+			},
+		},
+		{
+			spec: specutil.NewCol("string_varchar", "string", specutil.LitAttr("size", "255")),
+			expected: &schema.StringType{
+				T:    "varchar",
+				Size: 255,
+			},
+		},
+		{
+			spec: specutil.NewCol("string_test", "string", specutil.LitAttr("size", "10485761")),
+			expected: &schema.StringType{
+				T:    "text",
+				Size: 10_485_761,
+			},
+		},
+		{
+			spec: specutil.NewCol("varchar(255)", "varchar(255)"),
+			expected: &schema.StringType{
+				T:    "varchar",
+				Size: 255,
+			},
+		},
+		{
+			spec: specutil.NewCol("decimal(10, 2)", "decimal(10, 2)"),
+			expected: &schema.DecimalType{
+				T:         "decimal",
+				Scale:     2,
+				Precision: 10,
+			},
+		},
+		{
+			spec:     specutil.NewCol("enum", "enum", specutil.ListAttr("values", `"a"`, `"b"`, `"c"`)),
+			expected: &schema.EnumType{Values: []string{"a", "b", "c"}},
+		},
+		{
+			spec:     specutil.NewCol("bool", "boolean"),
+			expected: &schema.BoolType{T: "boolean"},
+		},
+		{
+			spec:     specutil.NewCol("decimal", "decimal", specutil.LitAttr("precision", "10"), specutil.LitAttr("scale", "2")),
+			expected: &schema.DecimalType{T: "decimal", Precision: 10, Scale: 2},
+		},
+		{
+			spec:     specutil.NewCol("float", "float", specutil.LitAttr("precision", "10")),
+			expected: &schema.FloatType{T: "real", Precision: 10},
+		},
+		{
+			spec:     specutil.NewCol("float", "float", specutil.LitAttr("precision", "25")),
+			expected: &schema.FloatType{T: "double precision", Precision: 25},
+		},
+		{
+			spec:     specutil.NewCol("cidr", "cidr"),
+			expected: &NetworkType{T: "cidr"},
+		},
+		{
+			spec:     specutil.NewCol("money", "money"),
+			expected: &CurrencyType{T: "money"},
+		},
+		{
+			spec:     specutil.NewCol("bit", "bit"),
+			expected: &BitType{T: "bit", Len: 1},
+		},
+		{
+			spec:     specutil.NewCol("bitvar", "bit varying"),
+			expected: &BitType{T: "bit varying"},
+		},
+		{
+			spec:     specutil.NewCol("bitvar8", "bit varying(8)"),
+			expected: &BitType{T: "bit varying", Len: 8},
+		},
+		{
+			spec:     specutil.NewCol("bit8", "bit(8)"),
+			expected: &BitType{T: "bit", Len: 8},
+		},
+	} {
+		t.Run(tt.spec.Name, func(t *testing.T) {
+			var s schema.Schema
+			err := UnmarshalSpec(hcl(tt.spec), schemahcl.Unmarshal, &s)
+			require.NoError(t, err)
+			tbl, ok := s.Table("table")
+			require.True(t, ok)
+			col, ok := tbl.Column(tt.spec.Name)
+			require.True(t, ok)
+			require.EqualValues(t, tt.expected, col.Type.Type)
+		})
+	}
+}
+
+func TestNotSupportedUnmarshalSpecColumnTypes(t *testing.T) {
+	for _, tt := range []struct {
+		spec        *sqlspec.Column
+		expectedErr string
+	}{
+		{
+			spec:        specutil.NewCol("uint", "uint"),
+			expectedErr: "postgres: failed converting to *schema.Schema: unsigned integers currently not supported",
+		},
+		{
+			spec:        specutil.NewCol("int8", "int8"),
+			expectedErr: "postgres: failed converting to *schema.Schema: 8-bit integers not supported",
+		},
+
+		{
+			spec:        specutil.NewCol("uint64", "uint64"),
+			expectedErr: "postgres: failed converting to *schema.Schema: unsigned integers currently not supported",
+		},
+	} {
+		t.Run(tt.spec.Name, func(t *testing.T) {
+			var s schema.Schema
+			err := UnmarshalSpec(hcl(tt.spec), schemahcl.Unmarshal, &s)
+			require.Equal(t, tt.expectedErr, err.Error())
+		})
+	}
+}
+
+// hcl returns an Atlas HCL document containing the column spec.
+func hcl(c *sqlspec.Column) []byte {
+	buf, err := schemahcl.Marshal(c)
+	if err != nil {
+		log.Fatalln(err)
+	}
+	tmpl := `
+schema "default" {
+}
+table "table" {
+	schema = "default"
+	%s
+}
+`
+	body := fmt.Sprintf(tmpl, buf)
+	return []byte(body)
+}