From 33360ab479cc29dd23adb848504bad01b7ff25b3 Mon Sep 17 00:00:00 2001 From: Jack Christensen Date: Sat, 16 Mar 2024 12:33:55 -0500 Subject: [PATCH] Add ScanLocation to pgtype.TimestamptzCodec If ScanLocation is set, it will be used to convert the time to the given location when scanning from the database. The Codec interface is now implemented by *pgtype.TimestamptzCodec instead of pgtype.TimestamptzCodec. This is technically a breaking change, but it is extremely unlikely that anyone is depending on this, and if there is downstream breakage it is trivial to fix. https://github.com/jackc/pgx/issues/1195 https://github.com/jackc/pgx/issues/1945 --- pgtype/pgtype_default.go | 2 +- pgtype/timestamptz.go | 39 ++++++++++++++++++++++++-------------- pgtype/timestamptz_test.go | 34 +++++++++++++++++++++++++++++++++ 3 files changed, 60 insertions(+), 15 deletions(-) diff --git a/pgtype/pgtype_default.go b/pgtype/pgtype_default.go index d56d95bd..c5f2b3ce 100644 --- a/pgtype/pgtype_default.go +++ b/pgtype/pgtype_default.go @@ -83,7 +83,7 @@ func initDefaultMap() { defaultMap.RegisterType(&Type{Name: "tid", OID: TIDOID, Codec: TIDCodec{}}) defaultMap.RegisterType(&Type{Name: "time", OID: TimeOID, Codec: TimeCodec{}}) defaultMap.RegisterType(&Type{Name: "timestamp", OID: TimestampOID, Codec: TimestampCodec{}}) - defaultMap.RegisterType(&Type{Name: "timestamptz", OID: TimestamptzOID, Codec: TimestamptzCodec{}}) + defaultMap.RegisterType(&Type{Name: "timestamptz", OID: TimestamptzOID, Codec: &TimestamptzCodec{}}) defaultMap.RegisterType(&Type{Name: "unknown", OID: UnknownOID, Codec: TextCodec{}}) defaultMap.RegisterType(&Type{Name: "uuid", OID: UUIDOID, Codec: UUIDCodec{}}) defaultMap.RegisterType(&Type{Name: "varbit", OID: VarbitOID, Codec: BitsCodec{}}) diff --git a/pgtype/timestamptz.go b/pgtype/timestamptz.go index f568fe30..7efbcffd 100644 --- a/pgtype/timestamptz.go +++ b/pgtype/timestamptz.go @@ -54,7 +54,7 @@ func (tstz *Timestamptz) Scan(src any) error { switch src := src.(type) { case string: - return scanPlanTextTimestamptzToTimestamptzScanner{}.Scan([]byte(src), tstz) + return (&scanPlanTextTimestamptzToTimestamptzScanner{}).Scan([]byte(src), tstz) case time.Time: *tstz = Timestamptz{Time: src, Valid: true} return nil @@ -124,17 +124,21 @@ func (tstz *Timestamptz) UnmarshalJSON(b []byte) error { return nil } -type TimestamptzCodec struct{} +type TimestamptzCodec struct { + // ScanLocation is the location to return scanned timestamptz values in. This does not change the instant in time that + // the timestamptz represents. + ScanLocation *time.Location +} -func (TimestamptzCodec) FormatSupported(format int16) bool { +func (*TimestamptzCodec) FormatSupported(format int16) bool { return format == TextFormatCode || format == BinaryFormatCode } -func (TimestamptzCodec) PreferredFormat() int16 { +func (*TimestamptzCodec) PreferredFormat() int16 { return BinaryFormatCode } -func (TimestamptzCodec) PlanEncode(m *Map, oid uint32, format int16, value any) EncodePlan { +func (*TimestamptzCodec) PlanEncode(m *Map, oid uint32, format int16, value any) EncodePlan { if _, ok := value.(TimestamptzValuer); !ok { return nil } @@ -220,27 +224,27 @@ func (encodePlanTimestamptzCodecText) Encode(value any, buf []byte) (newBuf []by return buf, nil } -func (TimestamptzCodec) PlanScan(m *Map, oid uint32, format int16, target any) ScanPlan { +func (c *TimestamptzCodec) PlanScan(m *Map, oid uint32, format int16, target any) ScanPlan { switch format { case BinaryFormatCode: switch target.(type) { case TimestamptzScanner: - return scanPlanBinaryTimestamptzToTimestamptzScanner{} + return &scanPlanBinaryTimestamptzToTimestamptzScanner{location: c.ScanLocation} } case TextFormatCode: switch target.(type) { case TimestamptzScanner: - return scanPlanTextTimestamptzToTimestamptzScanner{} + return &scanPlanTextTimestamptzToTimestamptzScanner{location: c.ScanLocation} } } return nil } -type scanPlanBinaryTimestamptzToTimestamptzScanner struct{} +type scanPlanBinaryTimestamptzToTimestamptzScanner struct{ location *time.Location } -func (scanPlanBinaryTimestamptzToTimestamptzScanner) Scan(src []byte, dst any) error { +func (plan *scanPlanBinaryTimestamptzToTimestamptzScanner) Scan(src []byte, dst any) error { scanner := (dst).(TimestamptzScanner) if src == nil { @@ -264,15 +268,18 @@ func (scanPlanBinaryTimestamptzToTimestamptzScanner) Scan(src []byte, dst any) e microsecFromUnixEpochToY2K/1000000+microsecSinceY2K/1000000, (microsecFromUnixEpochToY2K%1000000*1000)+(microsecSinceY2K%1000000*1000), ) + if plan.location != nil { + tim = tim.In(plan.location) + } tstz = Timestamptz{Time: tim, Valid: true} } return scanner.ScanTimestamptz(tstz) } -type scanPlanTextTimestamptzToTimestamptzScanner struct{} +type scanPlanTextTimestamptzToTimestamptzScanner struct{ location *time.Location } -func (scanPlanTextTimestamptzToTimestamptzScanner) Scan(src []byte, dst any) error { +func (plan *scanPlanTextTimestamptzToTimestamptzScanner) Scan(src []byte, dst any) error { scanner := (dst).(TimestamptzScanner) if src == nil { @@ -312,13 +319,17 @@ func (scanPlanTextTimestamptzToTimestamptzScanner) Scan(src []byte, dst any) err tim = time.Date(year, tim.Month(), tim.Day(), tim.Hour(), tim.Minute(), tim.Second(), tim.Nanosecond(), tim.Location()) } + if plan.location != nil { + tim = tim.In(plan.location) + } + tstz = Timestamptz{Time: tim, Valid: true} } return scanner.ScanTimestamptz(tstz) } -func (c TimestamptzCodec) DecodeDatabaseSQLValue(m *Map, oid uint32, format int16, src []byte) (driver.Value, error) { +func (c *TimestamptzCodec) DecodeDatabaseSQLValue(m *Map, oid uint32, format int16, src []byte) (driver.Value, error) { if src == nil { return nil, nil } @@ -336,7 +347,7 @@ func (c TimestamptzCodec) DecodeDatabaseSQLValue(m *Map, oid uint32, format int1 return tstz.Time, nil } -func (c TimestamptzCodec) DecodeValue(m *Map, oid uint32, format int16, src []byte) (any, error) { +func (c *TimestamptzCodec) DecodeValue(m *Map, oid uint32, format int16, src []byte) (any, error) { if src == nil { return nil, nil } diff --git a/pgtype/timestamptz_test.go b/pgtype/timestamptz_test.go index 0486ecdb..57248195 100644 --- a/pgtype/timestamptz_test.go +++ b/pgtype/timestamptz_test.go @@ -38,6 +38,40 @@ func TestTimestamptzCodec(t *testing.T) { }) } +func TestTimestamptzCodecWithLocationUTC(t *testing.T) { + skipCockroachDB(t, "Server does not support infinite timestamps (see https://github.com/cockroachdb/cockroach/issues/41564)") + + connTestRunner := defaultConnTestRunner + connTestRunner.AfterConnect = func(ctx context.Context, t testing.TB, conn *pgx.Conn) { + conn.TypeMap().RegisterType(&pgtype.Type{ + Name: "timestamptz", + OID: pgtype.TimestamptzOID, + Codec: &pgtype.TimestamptzCodec{ScanLocation: time.UTC}, + }) + } + + pgxtest.RunValueRoundTripTests(context.Background(), t, connTestRunner, nil, "timestamptz", []pgxtest.ValueRoundTripTest{ + {time.Date(2000, 1, 1, 0, 0, 0, 0, time.UTC), new(time.Time), isExpectedEq(time.Date(2000, 1, 1, 0, 0, 0, 0, time.UTC))}, + }) +} + +func TestTimestamptzCodecWithLocationLocal(t *testing.T) { + skipCockroachDB(t, "Server does not support infinite timestamps (see https://github.com/cockroachdb/cockroach/issues/41564)") + + connTestRunner := defaultConnTestRunner + connTestRunner.AfterConnect = func(ctx context.Context, t testing.TB, conn *pgx.Conn) { + conn.TypeMap().RegisterType(&pgtype.Type{ + Name: "timestamptz", + OID: pgtype.TimestamptzOID, + Codec: &pgtype.TimestamptzCodec{ScanLocation: time.Local}, + }) + } + + pgxtest.RunValueRoundTripTests(context.Background(), t, connTestRunner, nil, "timestamptz", []pgxtest.ValueRoundTripTest{ + {time.Date(2000, 1, 1, 0, 0, 0, 0, time.Local), new(time.Time), isExpectedEq(time.Date(2000, 1, 1, 0, 0, 0, 0, time.Local))}, + }) +} + // https://github.com/jackc/pgx/v4/pgtype/pull/128 func TestTimestamptzTranscodeBigTimeBinary(t *testing.T) { defaultConnTestRunner.RunTest(context.Background(), t, func(ctx context.Context, t testing.TB, conn *pgx.Conn) {