Repository: ariga/atlas
Pull request #307 (2021-12-13 14:04:28)

sql/postgres: postgres type registry


Patch:
diff --git a/schema/schemaspec/schemahcl/hcl.go b/schema/schemaspec/schemahcl/hcl.go
index 50929a9e..577d39e8 100644
--- a/schema/schemaspec/schemahcl/hcl.go
+++ b/schema/schemaspec/schemahcl/hcl.go
@@ -293,7 +293,7 @@ func (s *state) findTypeSpec(t string) (*schemaspec.TypeSpec, bool) {
 
 func hclType(spec *schemaspec.TypeSpec, typ *schemaspec.Type) (string, error) {
 	if len(typeFuncArgs(spec)) == 0 {
-		return typ.T, nil
+		return spec.Name, nil
 	}
 	args := make([]string, 0, len(spec.Attributes))
 	for _, param := range typeFuncArgs(spec) {
@@ -314,7 +314,7 @@ func hclType(spec *schemaspec.TypeSpec, typ *schemaspec.Type) (string, error) {
 			}
 		}
 	}
-	return fmt.Sprintf("%s(%s)", typ.T, strings.Join(args, ",")), nil
+	return fmt.Sprintf("%s(%s)", spec.Name, strings.Join(args, ",")), nil
 }
 
 func findAttr(attrs []*schemaspec.Attr, k string) (*schemaspec.Attr, bool) {
diff --git a/sql/internal/specutil/types.go b/sql/internal/specutil/types.go
index e517d6b0..ae14e541 100644
--- a/sql/internal/specutil/types.go
+++ b/sql/internal/specutil/types.go
@@ -78,8 +78,8 @@ func (r *TypeRegistry) Register(specs ...*schemaspec.TypeSpec) error {
 		if _, exists := r.FindByName(s.Name); exists {
 			return fmt.Errorf("specutil: type with name of %q already registered", s.T)
 		}
+		r.r = append(r.r, s)
 	}
-	r.r = append(r.r, specs...)
 	return nil
 }
 
@@ -191,9 +191,14 @@ func (r *TypeRegistry) Type(typ *schemaspec.Type, extra []*schemaspec.Attr, pars
 
 // TypeSpec returns a TypeSpec with the provided name.
 func TypeSpec(name string, attrs ...*schemaspec.TypeAttr) *schemaspec.TypeSpec {
+	return AliasTypeSpec(name, name, attrs...)
+}
+
+// AliasTypeSpec returns a TypeSpec with the provided name.
+func AliasTypeSpec(name, dbType string, attrs ...*schemaspec.TypeAttr) *schemaspec.TypeSpec {
 	return &schemaspec.TypeSpec{
 		Name:       name,
-		T:          name,
+		T:          dbType,
 		Attributes: attrs,
 	}
 }
diff --git a/sql/postgres/convert.go b/sql/postgres/convert.go
index 6bbbabe1..dc5354e9 100644
--- a/sql/postgres/convert.go
+++ b/sql/postgres/convert.go
@@ -174,7 +174,7 @@ func parseColumn(s string) (*columnDesc, error) {
 		}
 	)
 	switch c.parts[0] {
-	case tVarChar:
+	case tVarChar, tCharVar, tChar, tCharacter:
 		if len(c.parts) > 1 {
 			c.size, err = strconv.ParseInt(c.parts[1], 10, 64)
 			if err != nil {
@@ -202,6 +202,8 @@ func parseColumn(s string) (*columnDesc, error) {
 		c.precision = 53
 	case tReal, tFloat4:
 		c.precision = 24
+	default:
+		c.typ = s
 	}
 	return c, nil
 }
diff --git a/sql/postgres/sqlspec.go b/sql/postgres/sqlspec.go
index cbf3fd09..91776111 100644
--- a/sql/postgres/sqlspec.go
+++ b/sql/postgres/sqlspec.go
@@ -3,6 +3,7 @@ package postgres
 import (
 	"errors"
 	"fmt"
+	"reflect"
 	"strconv"
 	"strings"
 	"unicode"
@@ -102,7 +103,7 @@ func convertColumnType(spec *sqlspec.Column) (schema.Type, error) {
 	case sqlspec.TypeBoolean:
 		return &schema.BoolType{T: tBoolean}, nil
 	default:
-		return parseRawType(spec)
+		return parseRawType(spec.Type)
 	}
 }
 
@@ -196,16 +197,17 @@ func convertFloat(spec *sqlspec.Column) (schema.Type, error) {
 	return ft, nil
 }
 
-func parseRawType(spec *sqlspec.Column) (schema.Type, error) {
-	d, err := parseColumn(spec.Type)
+func parseRawType(typ string) (schema.Type, error) {
+	d, err := parseColumn(typ)
 	if err != nil {
 		return nil, err
 	}
 	// Normalize PostgreSQL array data types from "CREATE TABLE" format to
 	// "INFORMATION_SCHEMA" format (i.e. as it is inspected from the database).
-	if t, ok := arrayType(spec.Type); ok {
+	if t, ok := arrayType(typ); ok {
 		d = &columnDesc{typ: tArray, udt: t}
 	}
+
 	t := columnType(d)
 	// If the type is unknown (to us), we fallback to user-defined but expect
 	// to improve this in future versions by ensuring this against the database.
@@ -364,3 +366,58 @@ func arrayType(t string) (string, bool) {
 	}
 	return t[:strings.IndexByte(t, '[')], true
 }
+
+// TypeRegistry contains the supported TypeSpecs for the postgres driver.
+var TypeRegistry = specutil.NewRegistry(
+	specutil.TypeSpec(tBit, specutil.SizeTypeAttr(true)),
+	specutil.AliasTypeSpec("bit_varying", tBitVar, specutil.SizeTypeAttr(true)),
+	specutil.TypeSpec(tVarChar, specutil.SizeTypeAttr(true)),
+	specutil.TypeSpec(tChar, specutil.SizeTypeAttr(true)),
+	specutil.TypeSpec(tCharacter, specutil.SizeTypeAttr(true)),
+	specutil.TypeSpec(tInt2),
+	specutil.TypeSpec(tInt4),
+	specutil.TypeSpec(tInt8),
+	specutil.TypeSpec(tInt),
+	specutil.TypeSpec(tInteger),
+	specutil.TypeSpec(tSmallInt),
+	specutil.TypeSpec(tBigInt),
+	specutil.TypeSpec(tText),
+	specutil.TypeSpec(tBoolean),
+	specutil.TypeSpec(tBool),
+	specutil.TypeSpec(tBytea),
+	specutil.TypeSpec(tCIDR),
+	specutil.TypeSpec(tInet),
+	specutil.TypeSpec(tMACAddr),
+	specutil.TypeSpec(tMACAddr8),
+	specutil.TypeSpec(tCircle),
+	specutil.TypeSpec(tLine),
+	specutil.TypeSpec(tLseg),
+	specutil.TypeSpec(tBox),
+	specutil.TypeSpec(tPath),
+	specutil.TypeSpec(tPoint),
+	specutil.TypeSpec(tDate),
+	specutil.TypeSpec(tTime),
+	specutil.AliasTypeSpec("time_with_time_zone", tTimeWTZ),
+	specutil.AliasTypeSpec("time_without_time_zone", tTimeWOTZ),
+	specutil.TypeSpec(tTimestamp),
+	specutil.AliasTypeSpec("timestamp_with_time_zone", tTimestampWTZ),
+	specutil.AliasTypeSpec("timestamp_without_time_zone", tTimestampWOTZ),
+	specutil.TypeSpec("enum", &schemaspec.TypeAttr{Name: "values", Kind: reflect.Slice, Required: true}),
+	specutil.AliasTypeSpec("double_precision", tDouble),
+	specutil.TypeSpec(tReal),
+	specutil.TypeSpec(tFloat8),
+	specutil.TypeSpec(tFloat4),
+	specutil.TypeSpec(tNumeric),
+	specutil.TypeSpec(tDecimal),
+	specutil.TypeSpec(tSmallSerial),
+	specutil.TypeSpec(tSerial),
+	specutil.TypeSpec(tBigSerial),
+	specutil.TypeSpec(tSerial2),
+	specutil.TypeSpec(tSerial4),
+	specutil.TypeSpec(tSerial8),
+	specutil.TypeSpec(tXML),
+	specutil.TypeSpec(tJSON),
+	specutil.TypeSpec(tJSONB),
+	specutil.TypeSpec(tUUID),
+	specutil.TypeSpec(tMoney),
+)


Test patch:
diff --git a/sql/postgres/sqlspec_test.go b/sql/postgres/sqlspec_test.go
index 20203dd4..76105766 100644
--- a/sql/postgres/sqlspec_test.go
+++ b/sql/postgres/sqlspec_test.go
@@ -5,14 +5,18 @@ import (
 	"log"
 	"testing"
 
+	"ariga.io/atlas/schema/schemaspec"
 	"ariga.io/atlas/schema/schemaspec/schemahcl"
 	"ariga.io/atlas/sql/internal/specutil"
 	"ariga.io/atlas/sql/schema"
 	"ariga.io/atlas/sql/sqlspec"
+	"github.com/hashicorp/hcl/v2/hclwrite"
 
 	"github.com/stretchr/testify/require"
 )
 
+var hclState = schemahcl.New(schemahcl.WithTypes(TypeRegistry.Specs()))
+
 func TestSQLSpec(t *testing.T) {
 	f := `
 schema "schema" {
@@ -508,3 +512,238 @@ func TestMarshalSpecColumnType(t *testing.T) {
 		})
 	}
 }
+
+func TestTypes(t *testing.T) {
+	// TODO(rotemtam): enum, timestamptz, interval
+	for _, tt := range []struct {
+		typeExpr string
+		expected schema.Type
+	}{
+		{
+			typeExpr: "bit(10)",
+			expected: &BitType{T: tBit, Len: 10},
+		},
+		{
+			typeExpr: "bit_varying(10)",
+			expected: &BitType{T: tBitVar, Len: 10},
+		},
+		{
+			typeExpr: "boolean",
+			expected: &schema.BoolType{T: tBoolean},
+		},
+		{
+			typeExpr: "bool",
+			expected: &schema.BoolType{T: tBool},
+		},
+		{
+			typeExpr: "bytea",
+			expected: &schema.BinaryType{T: tBytea},
+		},
+		{
+			typeExpr: "varchar(255)",
+			expected: &schema.StringType{T: tVarChar, Size: 255},
+		},
+		{
+			typeExpr: "char(255)",
+			expected: &schema.StringType{T: tChar, Size: 255},
+		},
+		{
+			typeExpr: "character(255)",
+			expected: &schema.StringType{T: tCharacter, Size: 255},
+		},
+		{
+			typeExpr: "text",
+			expected: &schema.StringType{T: tText},
+		},
+		{
+			typeExpr: "smallint",
+			expected: &schema.IntegerType{T: tSmallInt},
+		},
+		{
+			typeExpr: "integer",
+			expected: &schema.IntegerType{T: tInteger},
+		},
+		{
+			typeExpr: "bigint",
+			expected: &schema.IntegerType{T: tBigInt},
+		},
+		{
+			typeExpr: "int",
+			expected: &schema.IntegerType{T: tInt},
+		},
+		{
+			typeExpr: "int2",
+			expected: &schema.IntegerType{T: tInt2},
+		},
+		{
+			typeExpr: "int4",
+			expected: &schema.IntegerType{T: tInt4},
+		},
+		{
+			typeExpr: "int8",
+			expected: &schema.IntegerType{T: tInt8},
+		},
+		{
+			typeExpr: "cidr",
+			expected: &NetworkType{T: tCIDR},
+		},
+		{
+			typeExpr: "inet",
+			expected: &NetworkType{T: tInet},
+		},
+		{
+			typeExpr: "macaddr",
+			expected: &NetworkType{T: tMACAddr},
+		},
+		{
+			typeExpr: "macaddr8",
+			expected: &NetworkType{T: tMACAddr8},
+		},
+		{
+			typeExpr: "circle",
+			expected: &schema.SpatialType{T: tCircle},
+		},
+		{
+			typeExpr: "line",
+			expected: &schema.SpatialType{T: tLine},
+		},
+		{
+			typeExpr: "lseg",
+			expected: &schema.SpatialType{T: tLseg},
+		},
+		{
+			typeExpr: "box",
+			expected: &schema.SpatialType{T: tBox},
+		},
+		{
+			typeExpr: "path",
+			expected: &schema.SpatialType{T: tPath},
+		},
+		{
+			typeExpr: "point",
+			expected: &schema.SpatialType{T: tPoint},
+		},
+
+		{
+			typeExpr: "date",
+			expected: &schema.TimeType{T: tDate},
+		},
+		{
+			typeExpr: "time",
+			expected: &schema.TimeType{T: tTime},
+		},
+		{
+			typeExpr: "time_with_time_zone",
+			expected: &schema.TimeType{T: tTimeWTZ},
+		},
+		{
+			typeExpr: "time_without_time_zone",
+			expected: &schema.TimeType{T: tTimeWOTZ},
+		},
+		{
+			typeExpr: "timestamp",
+			expected: &schema.TimeType{T: tTimestamp},
+		},
+		{
+			typeExpr: "timestamp_with_time_zone",
+			expected: &schema.TimeType{T: tTimestampWTZ},
+		},
+		{
+			typeExpr: "timestamp_without_time_zone",
+			expected: &schema.TimeType{T: tTimestampWOTZ},
+		},
+		{
+			typeExpr: "time",
+			expected: &schema.TimeType{T: tTime},
+		},
+		{
+			typeExpr: "real",
+			expected: &schema.FloatType{T: tReal, Precision: 24},
+		},
+		{
+			typeExpr: "float8",
+			expected: &schema.FloatType{T: tFloat8, Precision: 53},
+		},
+		{
+			typeExpr: "float4",
+			expected: &schema.FloatType{T: tFloat4, Precision: 24},
+		},
+		{
+			typeExpr: "numeric",
+			expected: &schema.DecimalType{T: tNumeric},
+		},
+		{
+			typeExpr: "decimal",
+			expected: &schema.DecimalType{T: tDecimal},
+		},
+		{
+			typeExpr: "smallserial",
+			expected: &SerialType{T: tSmallSerial},
+		},
+		{
+			typeExpr: "serial",
+			expected: &SerialType{T: tSerial},
+		},
+		{
+			typeExpr: "bigserial",
+			expected: &SerialType{T: tBigSerial},
+		},
+		{
+			typeExpr: "serial2",
+			expected: &SerialType{T: tSerial2},
+		},
+		{
+			typeExpr: "serial4",
+			expected: &SerialType{T: tSerial4},
+		},
+		{
+			typeExpr: "serial8",
+			expected: &SerialType{T: tSerial8},
+		},
+
+		{
+			typeExpr: "xml",
+			expected: &XMLType{T: tXML},
+		},
+		{
+			typeExpr: "json",
+			expected: &schema.JSONType{T: tJSON},
+		},
+		{
+			typeExpr: "jsonb",
+			expected: &schema.JSONType{T: tJSONB},
+		},
+		{
+			typeExpr: "uuid",
+			expected: &UUIDType{T: tUUID},
+		},
+		{
+			typeExpr: "money",
+			expected: &CurrencyType{T: tMoney},
+		},
+	} {
+		t.Run(tt.typeExpr, func(t *testing.T) {
+			// simulates sqlspec.Column until we change its Type field.
+			type col struct {
+				Type *schemaspec.Type `spec:"type"`
+				schemaspec.DefaultExtension
+			}
+			var test struct {
+				Columns []*col `spec:"column"`
+			}
+			doc := fmt.Sprintf(`column {
+	type = %s
+}
+`, tt.typeExpr)
+			err := hclState.UnmarshalSpec([]byte(doc), &test)
+			require.NoError(t, err)
+			column := test.Columns[0]
+			typ, err := TypeRegistry.Type(column.Type, column.Extra.Attrs, parseRawType)
+			require.NoError(t, err)
+			require.EqualValues(t, tt.expected, typ)
+			spec, err := hclState.MarshalSpec(&test)
+			require.NoError(t, err)
+			require.EqualValues(t, string(hclwrite.Format([]byte(doc))), string(hclwrite.Format(spec)))
+		})
+	}
+}