Add EnumType

enum_type.go

@@ -0,0 +1,163 @@
+package pgtype
+
+import errors "golang.org/x/xerrors"
+
+// EnumType represents an enum type. In the normal pgtype model a Go type maps to a PostgreSQL type and an instance
+// of a Go type maps to a PostgreSQL value of that type. EnumType is different in that an instance of EnumType
+// represents a PostgreSQL type. The zero value is not usable -- NewEnumType must be used as a constructor. In general,
+// an EnumType should not be used to represent a value. It should only be used as an encoder and decoder internal to
+// ConnInfo.
+type EnumType struct {
+	String string
+	Status Status
+
+	pgTypeName string            // PostgreSQL type name
+	members    []string          // enum members
+	membersMap map[string]string // map to quickly lookup member and reuse string instead of allocating
+}
+
+// NewEnumType initializes a new EnumType. It retains a read-only reference to members. members must not be changed.
+func NewEnumType(pgTypeName string, members []string) *EnumType {
+	et := &EnumType{pgTypeName: pgTypeName, members: members}
+	et.membersMap = make(map[string]string, len(members))
+	for _, m := range members {
+		et.membersMap[m] = m
+	}
+	return et
+}
+
+func (et *EnumType) CloneTypeValue() Value {
+	return &EnumType{
+		String: et.String,
+		Status: et.Status,
+
+		pgTypeName: et.pgTypeName,
+		members:    et.members,
+		membersMap: et.membersMap,
+	}
+}
+
+func (et *EnumType) PgTypeName() string {
+	return et.pgTypeName
+}
+
+func (et *EnumType) Members() []string {
+	return et.members
+}
+
+// Set assigns src to dst. Set purposely does not check that src is a member. This allows continued error free
+// operation in the event the PostgreSQL enum type is modified during a connection.
+func (dst *EnumType) Set(src interface{}) error {
+	if src == nil {
+		dst.Status = Null
+		return nil
+	}
+
+	if value, ok := src.(interface{ Get() interface{} }); ok {
+		value2 := value.Get()
+		if value2 != value {
+			return dst.Set(value2)
+		}
+	}
+
+	switch value := src.(type) {
+	case string:
+		dst.String = value
+		dst.Status = Present
+	case *string:
+		if value == nil {
+			dst.Status = Null
+		} else {
+			dst.String = *value
+			dst.Status = Present
+		}
+	case []byte:
+		if value == nil {
+			dst.Status = Null
+		} else {
+			dst.String = string(value)
+			dst.Status = Present
+		}
+	default:
+		if originalSrc, ok := underlyingStringType(src); ok {
+			return dst.Set(originalSrc)
+		}
+		return errors.Errorf("cannot convert %v to enum %s", value, dst.pgTypeName)
+	}
+
+	return nil
+}
+
+func (dst EnumType) Get() interface{} {
+	switch dst.Status {
+	case Present:
+		return dst.String
+	case Null:
+		return nil
+	default:
+		return dst.Status
+	}
+}
+
+func (src *EnumType) AssignTo(dst interface{}) error {
+	switch src.Status {
+	case Present:
+		switch v := dst.(type) {
+		case *string:
+			*v = src.String
+			return nil
+		case *[]byte:
+			*v = make([]byte, len(src.String))
+			copy(*v, src.String)
+			return nil
+		default:
+			if nextDst, retry := GetAssignToDstType(dst); retry {
+				return src.AssignTo(nextDst)
+			}
+			return errors.Errorf("unable to assign to %T", dst)
+		}
+	case Null:
+		return NullAssignTo(dst)
+	}
+
+	return errors.Errorf("cannot decode %#v into %T", src, dst)
+}
+
+func (dst *EnumType) DecodeText(ci *ConnInfo, src []byte) error {
+	if src == nil {
+		dst.Status = Null
+		return nil
+	}
+
+	// Lookup the string in membersMap to avoid an allocation.
+	if s, found := dst.membersMap[string(src)]; found {
+		dst.String = s
+	} else {
+		// If an enum type is modified after the initial connection it is possible to receive an unexpected value.
+		// Gracefully handle this situation. Purposely NOT modifying members and membersMap to allow for sharing members
+		// and membersMap between connections.
+		dst.String = string(src)
+	}
+	dst.Status = Present
+
+	return nil
+}
+
+func (dst *EnumType) DecodeBinary(ci *ConnInfo, src []byte) error {
+	return dst.DecodeText(ci, src)
+}
+
+func (src EnumType) EncodeText(ci *ConnInfo, buf []byte) ([]byte, error) {
+	switch src.Status {
+	case Null:
+		return nil, nil
+	case Undefined:
+		return nil, errUndefined
+	}
+
+	return append(buf, src.String...), nil
+}
+
+func (src EnumType) EncodeBinary(ci *ConnInfo, buf []byte) ([]byte, error) {
+	return src.EncodeText(ci, buf)
+}