From 7a3bc454e0ad5f080c323507c353fcad132ebd64 Mon Sep 17 00:00:00 2001 From: Jack Christensen Date: Thu, 20 Jan 2022 20:40:37 -0600 Subject: [PATCH] Convert TID to Codec --- pgtype/pgtype.go | 2 +- pgtype/tid.go | 233 +++++++++++++++++++++++++++++---------------- pgtype/tid_test.go | 66 +++---------- 3 files changed, 167 insertions(+), 134 deletions(-) diff --git a/pgtype/pgtype.go b/pgtype/pgtype.go index bcbb9f97..9b1ef595 100644 --- a/pgtype/pgtype.go +++ b/pgtype/pgtype.go @@ -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{}}) diff --git a/pgtype/tid.go b/pgtype/tid.go index 0108d219..624b3c2a 100644 --- a/pgtype/tid.go +++ b/pgtype/tid.go @@ -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 } diff --git a/pgtype/tid_test.go b/pgtype/tid_test.go index ef24005a..4203a471 100644 --- a/pgtype/tid_test.go +++ b/pgtype/tid_test.go @@ -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) - } - } -}