Repository: ariga/atlas
Pull request #422 (2022-01-16 15:55:53)

sql/postgres: export ParseType and improve array type parsing


Patch:
diff --git a/sql/postgres/convert.go b/sql/postgres/convert.go
index 710958c1..be045698 100644
--- a/sql/postgres/convert.go
+++ b/sql/postgres/convert.go
@@ -7,6 +7,7 @@ package postgres
 import (
 	"errors"
 	"fmt"
+	"regexp"
 	"strconv"
 	"strings"
 
@@ -146,6 +147,41 @@ func mustFormat(t schema.Type) string {
 	return s
 }
 
+// ParseType returns the schema.Type value represented by the given raw type.
+// The raw value is expected to follow the format in PostgreSQL information schema
+// or as an input for the CREATE TABLE statement.
+func ParseType(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(typ); ok {
+		d = &columnDesc{typ: TypeArray, 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.
+	if ut, ok := t.(*schema.UnsupportedType); ok {
+		t = &UserDefinedType{T: ut.T}
+	}
+	return t, nil
+}
+
+// reArray parses array declaration. See: https://postgresql.org/docs/current/arrays.html.
+var reArray = regexp.MustCompile(`(?i)(\w+)\s*(?:(?:\[\d*])+|\s+ARRAY\s*(?:\[\d*])*)`)
+
+// arrayType reports if the given string is an array type (e.g. int[], text[2]),
+// and returns its "udt_name" as it was inspected from the database.
+func arrayType(t string) (string, bool) {
+	matches := reArray.FindStringSubmatch(t)
+	if len(matches) != 2 {
+		return "", false
+	}
+	return matches[1], true
+}
+
 // columnDesc represents a column descriptor.
 type columnDesc struct {
 	typ       string
diff --git a/sql/postgres/sqlspec.go b/sql/postgres/sqlspec.go
index 542b748f..9169af75 100644
--- a/sql/postgres/sqlspec.go
+++ b/sql/postgres/sqlspec.go
@@ -3,8 +3,6 @@ package postgres
 import (
 	"reflect"
 	"strconv"
-	"strings"
-	"unicode"
 
 	"ariga.io/atlas/schema/schemaspec"
 	"ariga.io/atlas/schema/schemaspec/schemahcl"
@@ -61,26 +59,6 @@ func convertColumnType(spec *sqlspec.Column) (schema.Type, error) {
 	return TypeRegistry.Type(spec.Type, spec.Extra.Attrs)
 }
 
-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(typ); ok {
-		d = &columnDesc{typ: TypeArray, 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.
-	if ut, ok := t.(*schema.UnsupportedType); ok {
-		t = &UserDefinedType{T: ut.T}
-	}
-	return t, 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)
@@ -112,25 +90,10 @@ func columnTypeSpec(t schema.Type) (*sqlspec.Column, error) {
 	return &sqlspec.Column{Type: st}, nil
 }
 
-// arrayType reports if the given string is an array type (e.g. int[], text[2]),
-// and returns its "udt_name" as it was inspected from the database.
-func arrayType(t string) (string, bool) {
-	i, j := strings.LastIndexByte(t, '['), strings.LastIndexByte(t, ']')
-	if i == -1 || j == -1 {
-		return "", false
-	}
-	for _, r := range t[i+1 : j] {
-		if !unicode.IsDigit(r) {
-			return "", false
-		}
-	}
-	return t[:strings.IndexByte(t, '[')], true
-}
-
 // TypeRegistry contains the supported TypeSpecs for the Postgres driver.
 var TypeRegistry = specutil.NewRegistry(
 	specutil.WithFormatter(FormatType),
-	specutil.WithParser(parseRawType),
+	specutil.WithParser(ParseType),
 	specutil.WithSpecs(
 		specutil.TypeSpec(TypeBit, &schemaspec.TypeAttr{Name: "len", Kind: reflect.Int64}),
 		specutil.AliasTypeSpec("bit_varying", TypeBitVar, &schemaspec.TypeAttr{Name: "len", Kind: reflect.Int64}),


Test patch:
diff --git a/sql/postgres/sqlspec_test.go b/sql/postgres/sqlspec_test.go
index 5b79a070..a3df358f 100644
--- a/sql/postgres/sqlspec_test.go
+++ b/sql/postgres/sqlspec_test.go
@@ -401,6 +401,26 @@ func TestTypes(t *testing.T) {
 			typeExpr: `sql("int[]")`,
 			expected: &ArrayType{T: "int[]"},
 		},
+		{
+			typeExpr: `sql("int[2]")`,
+			expected: &ArrayType{T: "int[]"},
+		},
+		{
+			typeExpr: `sql("text[][]")`,
+			expected: &ArrayType{T: "text[]"},
+		},
+		{
+			typeExpr: `sql("integer [3][3]")`,
+			expected: &ArrayType{T: "integer[]"},
+		},
+		{
+			typeExpr: `sql("integer ARRAY[4]")`,
+			expected: &ArrayType{T: "integer[]"},
+		},
+		{
+			typeExpr: `sql("integer ARRAY")`,
+			expected: &ArrayType{T: "integer[]"},
+		},
 	} {
 		t.Run(tt.typeExpr, func(t *testing.T) {
 			var test schema.Schema