Convert TID to Codec

2022-01-20 20:40:37 -06:00
parent b10eb89fe4
commit 7a3bc454e0
3 changed files with 167 additions and 134 deletions
@@ -337,7 +337,7 @@ func NewConnInfo() *ConnInfo {
 	ci.RegisterDataType(DataType{Name: "polygon", OID: PolygonOID, Codec: PolygonCodec{}})
 	// ci.RegisterDataType(DataType{Value: &Record{}, Name: "record", OID: RecordOID})
 	ci.RegisterDataType(DataType{Name: "text", OID: TextOID, Codec: TextCodec{}})
-	ci.RegisterDataType(DataType{Value: &TID{}, Name: "tid", OID: TIDOID})
+	ci.RegisterDataType(DataType{Name: "tid", OID: TIDOID, Codec: TIDCodec{}})
 	ci.RegisterDataType(DataType{Value: &Time{}, Name: "time", OID: TimeOID})
 	ci.RegisterDataType(DataType{Name: "timestamp", OID: TimestampOID, Codec: TimestampCodec{}})
 	ci.RegisterDataType(DataType{Name: "timestamptz", OID: TimestamptzOID, Codec: TimestamptzCodec{}})
@@ -10,6 +10,14 @@ import (
 	"github.com/jackc/pgio"
 )

+type TIDScanner interface {
+	ScanTID(v TID) error
+}
+
+type TIDValuer interface {
+	TIDValue() (TID, error)
+}
+
 // TID is PostgreSQL's Tuple Identifier type.
 //
 // When one does
@@ -27,40 +35,148 @@ type TID struct {
 	Valid        bool
 }

-func (dst *TID) Set(src interface{}) error {
-	return fmt.Errorf("cannot convert %v to TID", src)
+func (b *TID) ScanTID(v TID) error {
+	*b = v
+	return nil
 }

-func (dst TID) Get() interface{} {
-	if !dst.Valid {
-		return nil
-	}
-	return dst
+func (b TID) TIDValue() (TID, error) {
+	return b, nil
 }

-func (src *TID) AssignTo(dst interface{}) error {
-	if !src.Valid {
-		return fmt.Errorf("cannot assign %v to %T", src, dst)
-	}
-
-	switch v := dst.(type) {
-	case *string:
-		*v = fmt.Sprintf(`(%d,%d)`, src.BlockNumber, src.OffsetNumber)
-		return nil
-	default:
-		if nextDst, retry := GetAssignToDstType(dst); retry {
-			return src.AssignTo(nextDst)
-		}
-		return fmt.Errorf("unable to assign to %T", dst)
-	}
-}
-
-func (dst *TID) DecodeText(ci *ConnInfo, src []byte) error {
+// Scan implements the database/sql Scanner interface.
+func (dst *TID) Scan(src interface{}) error {
 	if src == nil {
 		*dst = TID{}
 		return nil
 	}

+	switch src := src.(type) {
+	case string:
+		return scanPlanTextAnyToTIDScanner{}.Scan(nil, 0, TextFormatCode, []byte(src), dst)
+	}
+
+	return fmt.Errorf("cannot scan %T", src)
+}
+
+// Value implements the database/sql/driver Valuer interface.
+func (src TID) Value() (driver.Value, error) {
+	if !src.Valid {
+		return nil, nil
+	}
+
+	buf, err := TIDCodec{}.PlanEncode(nil, 0, TextFormatCode, src).Encode(src, nil)
+	if err != nil {
+		return nil, err
+	}
+	return string(buf), err
+}
+
+type TIDCodec struct{}
+
+func (TIDCodec) FormatSupported(format int16) bool {
+	return format == TextFormatCode || format == BinaryFormatCode
+}
+
+func (TIDCodec) PreferredFormat() int16 {
+	return BinaryFormatCode
+}
+
+func (TIDCodec) PlanEncode(ci *ConnInfo, oid uint32, format int16, value interface{}) EncodePlan {
+	if _, ok := value.(TIDValuer); !ok {
+		return nil
+	}
+
+	switch format {
+	case BinaryFormatCode:
+		return encodePlanTIDCodecBinary{}
+	case TextFormatCode:
+		return encodePlanTIDCodecText{}
+	}
+
+	return nil
+}
+
+type encodePlanTIDCodecBinary struct{}
+
+func (encodePlanTIDCodecBinary) Encode(value interface{}, buf []byte) (newBuf []byte, err error) {
+	tid, err := value.(TIDValuer).TIDValue()
+	if err != nil {
+		return nil, err
+	}
+
+	if !tid.Valid {
+		return nil, nil
+	}
+
+	buf = pgio.AppendUint32(buf, tid.BlockNumber)
+	buf = pgio.AppendUint16(buf, tid.OffsetNumber)
+	return buf, nil
+}
+
+type encodePlanTIDCodecText struct{}
+
+func (encodePlanTIDCodecText) Encode(value interface{}, buf []byte) (newBuf []byte, err error) {
+	tid, err := value.(TIDValuer).TIDValue()
+	if err != nil {
+		return nil, err
+	}
+
+	if !tid.Valid {
+		return nil, nil
+	}
+
+	buf = append(buf, fmt.Sprintf(`(%d,%d)`, tid.BlockNumber, tid.OffsetNumber)...)
+	return buf, nil
+}
+
+func (TIDCodec) PlanScan(ci *ConnInfo, oid uint32, format int16, target interface{}, actualTarget bool) ScanPlan {
+
+	switch format {
+	case BinaryFormatCode:
+		switch target.(type) {
+		case TIDScanner:
+			return scanPlanBinaryTIDToTIDScanner{}
+		}
+	case TextFormatCode:
+		switch target.(type) {
+		case TIDScanner:
+			return scanPlanTextAnyToTIDScanner{}
+		}
+	}
+
+	return nil
+}
+
+type scanPlanBinaryTIDToTIDScanner struct{}
+
+func (scanPlanBinaryTIDToTIDScanner) Scan(ci *ConnInfo, oid uint32, formatCode int16, src []byte, dst interface{}) error {
+	scanner := (dst).(TIDScanner)
+
+	if src == nil {
+		return scanner.ScanTID(TID{})
+	}
+
+	if len(src) != 6 {
+		return fmt.Errorf("invalid length for tid: %v", len(src))
+	}
+
+	return scanner.ScanTID(TID{
+		BlockNumber:  binary.BigEndian.Uint32(src),
+		OffsetNumber: binary.BigEndian.Uint16(src[4:]),
+		Valid:        true,
+	})
+}
+
+type scanPlanTextAnyToTIDScanner struct{}
+
+func (scanPlanTextAnyToTIDScanner) Scan(ci *ConnInfo, oid uint32, formatCode int16, src []byte, dst interface{}) error {
+	scanner := (dst).(TIDScanner)
+
+	if src == nil {
+		return scanner.ScanTID(TID{})
+	}
+
 	if len(src) < 5 {
 		return fmt.Errorf("invalid length for tid: %v", len(src))
 	}
@@ -80,67 +196,22 @@ func (dst *TID) DecodeText(ci *ConnInfo, src []byte) error {
 		return err
 	}

-	*dst = TID{BlockNumber: uint32(blockNumber), OffsetNumber: uint16(offsetNumber), Valid: true}
-	return nil
+	return scanner.ScanTID(TID{BlockNumber: uint32(blockNumber), OffsetNumber: uint16(offsetNumber), Valid: true})
 }

-func (dst *TID) DecodeBinary(ci *ConnInfo, src []byte) error {
+func (c TIDCodec) DecodeDatabaseSQLValue(ci *ConnInfo, oid uint32, format int16, src []byte) (driver.Value, error) {
+	return codecDecodeToTextFormat(c, ci, oid, format, src)
+}
+
+func (c TIDCodec) DecodeValue(ci *ConnInfo, oid uint32, format int16, src []byte) (interface{}, error) {
 	if src == nil {
-		*dst = TID{}
-		return nil
-	}
-
-	if len(src) != 6 {
-		return fmt.Errorf("invalid length for tid: %v", len(src))
-	}
-
-	*dst = TID{
-		BlockNumber:  binary.BigEndian.Uint32(src),
-		OffsetNumber: binary.BigEndian.Uint16(src[4:]),
-		Valid:        true,
-	}
-	return nil
-}
-
-func (src TID) EncodeText(ci *ConnInfo, buf []byte) ([]byte, error) {
-	if !src.Valid {
 		return nil, nil
 	}

-	buf = append(buf, fmt.Sprintf(`(%d,%d)`, src.BlockNumber, src.OffsetNumber)...)
-	return buf, nil
-}
-
-func (src TID) EncodeBinary(ci *ConnInfo, buf []byte) ([]byte, error) {
-	if !src.Valid {
-		return nil, nil
+	var tid TID
+	err := codecScan(c, ci, oid, format, src, &tid)
+	if err != nil {
+		return nil, err
 	}
-
-	buf = pgio.AppendUint32(buf, src.BlockNumber)
-	buf = pgio.AppendUint16(buf, src.OffsetNumber)
-	return buf, nil
-}
-
-// Scan implements the database/sql Scanner interface.
-func (dst *TID) Scan(src interface{}) error {
-	if src == nil {
-		*dst = TID{}
-		return 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 fmt.Errorf("cannot scan %T", src)
-}
-
-// Value implements the database/sql/driver Valuer interface.
-func (src TID) Value() (driver.Value, error) {
-	return EncodeValueText(src)
+	return tid, nil
 }
@@ -1,62 +1,24 @@
 package pgtype_test

 import (
-	"reflect"
 	"testing"

 	"github.com/jackc/pgx/v5/pgtype"
-	"github.com/jackc/pgx/v5/pgtype/testutil"
 )

-func TestTIDTranscode(t *testing.T) {
-	testutil.TestSuccessfulTranscode(t, "tid", []interface{}{
-		&pgtype.TID{BlockNumber: 42, OffsetNumber: 43, Valid: true},
-		&pgtype.TID{BlockNumber: 4294967295, OffsetNumber: 65535, Valid: true},
-		&pgtype.TID{},
+func TestTIDCodec(t *testing.T) {
+	testPgxCodec(t, "tid", []PgxTranscodeTestCase{
+		{
+			pgtype.TID{BlockNumber: 42, OffsetNumber: 43, Valid: true},
+			new(pgtype.TID),
+			isExpectedEq(pgtype.TID{BlockNumber: 42, OffsetNumber: 43, Valid: true}),
+		},
+		{
+			pgtype.TID{BlockNumber: 4294967295, OffsetNumber: 65535, Valid: true},
+			new(pgtype.TID),
+			isExpectedEq(pgtype.TID{BlockNumber: 4294967295, OffsetNumber: 65535, Valid: true}),
+		},
+		{pgtype.TID{}, new(pgtype.TID), isExpectedEq(pgtype.TID{})},
+		{nil, new(pgtype.TID), isExpectedEq(pgtype.TID{})},
 	})
 }
-
-func TestTIDAssignTo(t *testing.T) {
-	var s string
-	var sp *string
-
-	simpleTests := []struct {
-		src      pgtype.TID
-		dst      interface{}
-		expected interface{}
-	}{
-		{src: pgtype.TID{BlockNumber: 42, OffsetNumber: 43, Valid: true}, dst: &s, expected: "(42,43)"},
-		{src: pgtype.TID{BlockNumber: 4294967295, OffsetNumber: 65535, Valid: true}, dst: &s, expected: "(4294967295,65535)"},
-	}
-
-	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(); dst != tt.expected {
-			t.Errorf("%d: expected %v to assign %v, but result was %v", i, tt.src, tt.expected, dst)
-		}
-	}
-
-	pointerAllocTests := []struct {
-		src      pgtype.TID
-		dst      interface{}
-		expected interface{}
-	}{
-		{src: pgtype.TID{BlockNumber: 42, OffsetNumber: 43, Valid: true}, dst: &sp, expected: "(42,43)"},
-		{src: pgtype.TID{BlockNumber: 4294967295, OffsetNumber: 65535, Valid: true}, dst: &sp, expected: "(4294967295,65535)"},
-	}
-
-	for i, tt := range pointerAllocTests {
-		err := tt.src.AssignTo(tt.dst)
-		if err != nil {
-			t.Errorf("%d: %v", i, err)
-		}
-
-		if dst := reflect.ValueOf(tt.dst).Elem().Elem().Interface(); dst != tt.expected {
-			t.Errorf("%d: expected %v to assign %v, but result was %v", i, tt.src, tt.expected, dst)
-		}
-	}
-}