Add support for array of enum

fixes #338

enum_array.go

@@ -0,0 +1,212 @@
+package pgtype
+
+import (
+	"database/sql/driver"
+
+	"github.com/pkg/errors"
+)
+
+type EnumArray struct {
+	Elements   []GenericText
+	Dimensions []ArrayDimension
+	Status     Status
+}
+
+func (dst *EnumArray) Set(src interface{}) error {
+	// untyped nil and typed nil interfaces are different
+	if src == nil {
+		*dst = EnumArray{Status: Null}
+		return nil
+	}
+
+	switch value := src.(type) {
+
+	case []string:
+		if value == nil {
+			*dst = EnumArray{Status: Null}
+		} else if len(value) == 0 {
+			*dst = EnumArray{Status: Present}
+		} else {
+			elements := make([]GenericText, len(value))
+			for i := range value {
+				if err := elements[i].Set(value[i]); err != nil {
+					return err
+				}
+			}
+			*dst = EnumArray{
+				Elements:   elements,
+				Dimensions: []ArrayDimension{{Length: int32(len(elements)), LowerBound: 1}},
+				Status:     Present,
+			}
+		}
+
+	default:
+		if originalSrc, ok := underlyingSliceType(src); ok {
+			return dst.Set(originalSrc)
+		}
+		return errors.Errorf("cannot convert %v to EnumArray", value)
+	}
+
+	return nil
+}
+
+func (dst *EnumArray) Get() interface{} {
+	switch dst.Status {
+	case Present:
+		return dst
+	case Null:
+		return nil
+	default:
+		return dst.Status
+	}
+}
+
+func (src *EnumArray) AssignTo(dst interface{}) error {
+	switch src.Status {
+	case Present:
+		switch v := dst.(type) {
+
+		case *[]string:
+			*v = make([]string, len(src.Elements))
+			for i := range src.Elements {
+				if err := src.Elements[i].AssignTo(&((*v)[i])); err != nil {
+					return err
+				}
+			}
+			return nil
+
+		default:
+			if nextDst, retry := GetAssignToDstType(dst); retry {
+				return src.AssignTo(nextDst)
+			}
+		}
+	case Null:
+		return NullAssignTo(dst)
+	}
+
+	return errors.Errorf("cannot decode %v into %T", src, dst)
+}
+
+func (dst *EnumArray) DecodeText(ci *ConnInfo, src []byte) error {
+	if src == nil {
+		*dst = EnumArray{Status: Null}
+		return nil
+	}
+
+	uta, err := ParseUntypedTextArray(string(src))
+	if err != nil {
+		return err
+	}
+
+	var elements []GenericText
+
+	if len(uta.Elements) > 0 {
+		elements = make([]GenericText, len(uta.Elements))
+
+		for i, s := range uta.Elements {
+			var elem GenericText
+			var elemSrc []byte
+			if s != "NULL" {
+				elemSrc = []byte(s)
+			}
+			err = elem.DecodeText(ci, elemSrc)
+			if err != nil {
+				return err
+			}
+
+			elements[i] = elem
+		}
+	}
+
+	*dst = EnumArray{Elements: elements, Dimensions: uta.Dimensions, Status: Present}
+
+	return nil
+}
+
+func (src *EnumArray) EncodeText(ci *ConnInfo, buf []byte) ([]byte, error) {
+	switch src.Status {
+	case Null:
+		return nil, nil
+	case Undefined:
+		return nil, errUndefined
+	}
+
+	if len(src.Dimensions) == 0 {
+		return append(buf, '{', '}'), nil
+	}
+
+	buf = EncodeTextArrayDimensions(buf, src.Dimensions)
+
+	// dimElemCounts is the multiples of elements that each array lies on. For
+	// example, a single dimension array of length 4 would have a dimElemCounts of
+	// [4]. A multi-dimensional array of lengths [3,5,2] would have a
+	// dimElemCounts of [30,10,2]. This is used to simplify when to render a '{'
+	// or '}'.
+	dimElemCounts := make([]int, len(src.Dimensions))
+	dimElemCounts[len(src.Dimensions)-1] = int(src.Dimensions[len(src.Dimensions)-1].Length)
+	for i := len(src.Dimensions) - 2; i > -1; i-- {
+		dimElemCounts[i] = int(src.Dimensions[i].Length) * dimElemCounts[i+1]
+	}
+
+	inElemBuf := make([]byte, 0, 32)
+	for i, elem := range src.Elements {
+		if i > 0 {
+			buf = append(buf, ',')
+		}
+
+		for _, dec := range dimElemCounts {
+			if i%dec == 0 {
+				buf = append(buf, '{')
+			}
+		}
+
+		elemBuf, err := elem.EncodeText(ci, inElemBuf)
+		if err != nil {
+			return nil, err
+		}
+		if elemBuf == nil {
+			buf = append(buf, `NULL`...)
+		} else {
+			buf = append(buf, QuoteArrayElementIfNeeded(string(elemBuf))...)
+		}
+
+		for _, dec := range dimElemCounts {
+			if (i+1)%dec == 0 {
+				buf = append(buf, '}')
+			}
+		}
+	}
+
+	return buf, nil
+}
+
+// Scan implements the database/sql Scanner interface.
+func (dst *EnumArray) Scan(src interface{}) error {
+	if src == nil {
+		return dst.DecodeText(nil, nil)
+	}
+
+	switch src := src.(type) {
+	case string:
+		return dst.DecodeText(nil, []byte(src))
+	case []byte:
+		srcCopy := make([]byte, len(src))
+		copy(srcCopy, src)
+		return dst.DecodeText(nil, srcCopy)
+	}
+
+	return errors.Errorf("cannot scan %T", src)
+}
+
+// Value implements the database/sql/driver Valuer interface.
+func (src *EnumArray) Value() (driver.Value, error) {
+	buf, err := src.EncodeText(nil, nil)
+	if err != nil {
+		return nil, err
+	}
+	if buf == nil {
+		return nil, nil
+	}
+
+	return string(buf), nil
+}

enum_array_test.go

@@ -0,0 +1,150 @@
+package pgtype_test
+
+import (
+	"reflect"
+	"testing"
+
+	"github.com/jackc/pgx/pgtype"
+	"github.com/jackc/pgx/pgtype/testutil"
+)
+
+func TestEnumArrayTranscode(t *testing.T) {
+	setupConn := testutil.MustConnectPgx(t)
+	defer testutil.MustClose(t, setupConn)
+
+	if _, err := setupConn.Exec("drop type if exists color"); err != nil {
+		t.Fatal(err)
+	}
+	if _, err := setupConn.Exec("create type color as enum ('red', 'green', 'blue')"); err != nil {
+		t.Fatal(err)
+	}
+
+	testutil.TestSuccessfulTranscode(t, "color[]", []interface{}{
+		&pgtype.EnumArray{
+			Elements:   nil,
+			Dimensions: nil,
+			Status:     pgtype.Present,
+		},
+		&pgtype.EnumArray{
+			Elements: []pgtype.GenericText{
+				pgtype.GenericText{String: "red", Status: pgtype.Present},
+				pgtype.GenericText{Status: pgtype.Null},
+			},
+			Dimensions: []pgtype.ArrayDimension{{Length: 2, LowerBound: 1}},
+			Status:     pgtype.Present,
+		},
+		&pgtype.EnumArray{Status: pgtype.Null},
+		&pgtype.EnumArray{
+			Elements: []pgtype.GenericText{
+				pgtype.GenericText{String: "red", Status: pgtype.Present},
+				pgtype.GenericText{String: "green", Status: pgtype.Present},
+				pgtype.GenericText{String: "blue", Status: pgtype.Present},
+				pgtype.GenericText{String: "red", Status: pgtype.Present},
+			},
+			Dimensions: []pgtype.ArrayDimension{
+				{Length: 2, LowerBound: 4},
+				{Length: 2, LowerBound: 2},
+			},
+			Status: pgtype.Present,
+		},
+	})
+}
+
+func TestEnumArrayArraySet(t *testing.T) {
+	successfulTests := []struct {
+		source interface{}
+		result pgtype.EnumArray
+	}{
+		{
+			source: []string{"foo"},
+			result: pgtype.EnumArray{
+				Elements:   []pgtype.GenericText{{String: "foo", Status: pgtype.Present}},
+				Dimensions: []pgtype.ArrayDimension{{LowerBound: 1, Length: 1}},
+				Status:     pgtype.Present},
+		},
+		{
+			source: (([]string)(nil)),
+			result: pgtype.EnumArray{Status: pgtype.Null},
+		},
+	}
+
+	for i, tt := range successfulTests {
+		var r pgtype.EnumArray
+		err := r.Set(tt.source)
+		if err != nil {
+			t.Errorf("%d: %v", i, err)
+		}
+
+		if !reflect.DeepEqual(r, tt.result) {
+			t.Errorf("%d: expected %v to convert to %v, but it was %v", i, tt.source, tt.result, r)
+		}
+	}
+}
+
+func TestEnumArrayArrayAssignTo(t *testing.T) {
+	var stringSlice []string
+	type _stringSlice []string
+	var namedStringSlice _stringSlice
+
+	simpleTests := []struct {
+		src      pgtype.EnumArray
+		dst      interface{}
+		expected interface{}
+	}{
+		{
+			src: pgtype.EnumArray{
+				Elements:   []pgtype.GenericText{{String: "foo", Status: pgtype.Present}},
+				Dimensions: []pgtype.ArrayDimension{{LowerBound: 1, Length: 1}},
+				Status:     pgtype.Present,
+			},
+			dst:      &stringSlice,
+			expected: []string{"foo"},
+		},
+		{
+			src: pgtype.EnumArray{
+				Elements:   []pgtype.GenericText{{String: "bar", Status: pgtype.Present}},
+				Dimensions: []pgtype.ArrayDimension{{LowerBound: 1, Length: 1}},
+				Status:     pgtype.Present,
+			},
+			dst:      &namedStringSlice,
+			expected: _stringSlice{"bar"},
+		},
+		{
+			src:      pgtype.EnumArray{Status: pgtype.Null},
+			dst:      &stringSlice,
+			expected: (([]string)(nil)),
+		},
+	}
+
+	for i, tt := range simpleTests {
+		err := tt.src.AssignTo(tt.dst)
+		if err != nil {
+			t.Errorf("%d: %v", i, err)
+		}
+
+		if dst := reflect.ValueOf(tt.dst).Elem().Interface(); !reflect.DeepEqual(dst, tt.expected) {
+			t.Errorf("%d: expected %v to assign %v, but result was %v", i, tt.src, tt.expected, dst)
+		}
+	}
+
+	errorTests := []struct {
+		src pgtype.EnumArray
+		dst interface{}
+	}{
+		{
+			src: pgtype.EnumArray{
+				Elements:   []pgtype.GenericText{{Status: pgtype.Null}},
+				Dimensions: []pgtype.ArrayDimension{{LowerBound: 1, Length: 1}},
+				Status:     pgtype.Present,
+			},
+			dst: &stringSlice,
+		},
+	}
+
+	for i, tt := range errorTests {
+		err := tt.src.AssignTo(tt.dst)
+		if err == nil {
+			t.Errorf("%d: expected error but none was returned (%v -> %v)", i, tt.src, tt.dst)
+		}
+	}
+}

typed_array_gen.sh

@@ -16,4 +16,8 @@ erb pgtype_array_type=ACLItemArray pgtype_element_type=ACLItem go_array_types=[]
 erb pgtype_array_type=HstoreArray pgtype_element_type=Hstore go_array_types=[]map[string]string element_type_name=hstore text_null=NULL binary_format=true typed_array.go.erb > hstore_array.go
 erb pgtype_array_type=NumericArray pgtype_element_type=Numeric go_array_types=[]float32,[]float64 element_type_name=numeric text_null=NULL binary_format=true typed_array.go.erb > numeric_array.go
 erb pgtype_array_type=UUIDArray pgtype_element_type=UUID go_array_types=[][16]byte,[][]byte,[]string element_type_name=uuid text_null=NULL binary_format=true typed_array.go.erb > uuid_array.go
+
+# While the binary format is theoretically possible it is only practical to use the text format. In addition, the text format for NULL enums is unquoted so TextArray or a possible GenericTextArray cannot be used.
+erb pgtype_array_type=EnumArray pgtype_element_type=GenericText go_array_types=[]string text_null='NULL' binary_format=false typed_array.go.erb > enum_array.go
+
 goimports -w *_array.go