Repository: jackc/pgx
Pull request #2000 (2024-05-04 17:11:27)

Document that generic helpers call rows.Close()
Existing generic helpers always call `defer rows.Close()`. Examples of their usage also omit external `defer rows.Close()` call.

For clarity, state that explicitly, because that's another point why one would want to switch to generic helpers from manually written `rows.Next()` loop.


Patch:
diff --git a/doc.go b/doc.go
index db99fc4c..911571b1 100644
--- a/doc.go
+++ b/doc.go
@@ -23,7 +23,7 @@ github.com/jackc/pgx/v5/pgxpool for a concurrency safe connection pool.
 Query Interface
 
 pgx implements Query in the familiar database/sql style. However, pgx provides generic functions such as CollectRows and
-ForEachRow that are a simpler and safer way of processing rows than manually calling rows.Next(), rows.Scan, and
+ForEachRow that are a simpler and safer way of processing rows than manually calling defer rows.Close(), rows.Next(), rows.Scan, and
 rows.Err().
 
 CollectRows can be used collect all returned rows into a slice.
diff --git a/pgconn/config.go b/pgconn/config.go
index 13def1e8..95b7c912 100644
--- a/pgconn/config.go
+++ b/pgconn/config.go
@@ -19,6 +19,7 @@ import (
 
 	"github.com/jackc/pgpassfile"
 	"github.com/jackc/pgservicefile"
+	"github.com/jackc/pgx/v5/pgconn/ctxwatch"
 	"github.com/jackc/pgx/v5/pgproto3"
 )
 
@@ -39,7 +40,12 @@ type Config struct {
 	DialFunc       DialFunc   // e.g. net.Dialer.DialContext
 	LookupFunc     LookupFunc // e.g. net.Resolver.LookupHost
 	BuildFrontend  BuildFrontendFunc
-	RuntimeParams  map[string]string // Run-time parameters to set on connection as session default values (e.g. search_path or application_name)
+
+	// BuildContextWatcherHandler is called to create a ContextWatcherHandler for a connection. The handler is called
+	// when a context passed to a PgConn method is canceled.
+	BuildContextWatcherHandler func(*PgConn) ctxwatch.Handler
+
+	RuntimeParams map[string]string // Run-time parameters to set on connection as session default values (e.g. search_path or application_name)
 
 	KerberosSrvName string
 	KerberosSpn     string
@@ -266,6 +272,9 @@ func ParseConfigWithOptions(connString string, options ParseConfigOptions) (*Con
 		BuildFrontend: func(r io.Reader, w io.Writer) *pgproto3.Frontend {
 			return pgproto3.NewFrontend(r, w)
 		},
+		BuildContextWatcherHandler: func(pgConn *PgConn) ctxwatch.Handler {
+			return &DeadlineContextWatcherHandler{Conn: pgConn.conn}
+		},
 		OnPgError: func(_ *PgConn, pgErr *PgError) bool {
 			// we want to automatically close any fatal errors
 			if strings.EqualFold(pgErr.Severity, "FATAL") {
diff --git a/pgconn/internal/ctxwatch/context_watcher.go b/pgconn/ctxwatch/context_watcher.go
similarity index 71%
rename from pgconn/internal/ctxwatch/context_watcher.go
rename to pgconn/ctxwatch/context_watcher.go
index b39cb3ee..db8884eb 100644
--- a/pgconn/internal/ctxwatch/context_watcher.go
+++ b/pgconn/ctxwatch/context_watcher.go
@@ -8,9 +8,8 @@ import (
 // ContextWatcher watches a context and performs an action when the context is canceled. It can watch one context at a
 // time.
 type ContextWatcher struct {
-	onCancel             func()
-	onUnwatchAfterCancel func()
-	unwatchChan          chan struct{}
+	handler     Handler
+	unwatchChan chan struct{}
 
 	lock              sync.Mutex
 	watchInProgress   bool
@@ -20,11 +19,10 @@ type ContextWatcher struct {
 // NewContextWatcher returns a ContextWatcher. onCancel will be called when a watched context is canceled.
 // OnUnwatchAfterCancel will be called when Unwatch is called and the watched context had already been canceled and
 // onCancel called.
-func NewContextWatcher(onCancel func(), onUnwatchAfterCancel func()) *ContextWatcher {
+func NewContextWatcher(handler Handler) *ContextWatcher {
 	cw := &ContextWatcher{
-		onCancel:             onCancel,
-		onUnwatchAfterCancel: onUnwatchAfterCancel,
-		unwatchChan:          make(chan struct{}),
+		handler:     handler,
+		unwatchChan: make(chan struct{}),
 	}
 
 	return cw
@@ -46,7 +44,7 @@ func (cw *ContextWatcher) Watch(ctx context.Context) {
 		go func() {
 			select {
 			case <-ctx.Done():
-				cw.onCancel()
+				cw.handler.HandleCancel(ctx)
 				cw.onCancelWasCalled = true
 				<-cw.unwatchChan
 			case <-cw.unwatchChan:
@@ -66,8 +64,17 @@ func (cw *ContextWatcher) Unwatch() {
 	if cw.watchInProgress {
 		cw.unwatchChan <- struct{}{}
 		if cw.onCancelWasCalled {
-			cw.onUnwatchAfterCancel()
+			cw.handler.HandleUnwatchAfterCancel()
 		}
 		cw.watchInProgress = false
 	}
 }
+
+type Handler interface {
+	// HandleCancel is called when the context that a ContextWatcher is currently watching is canceled. canceledCtx is the
+	// context that was canceled.
+	HandleCancel(canceledCtx context.Context)
+
+	// HandleUnwatchAfterCancel is called when a ContextWatcher that called HandleCancel on this Handler is unwatched.
+	HandleUnwatchAfterCancel()
+}
diff --git a/pgconn/doc.go b/pgconn/doc.go
index e3242cf4..8cf8688a 100644
--- a/pgconn/doc.go
+++ b/pgconn/doc.go
@@ -20,13 +20,17 @@ result. The ReadAll method reads all query results into memory.
 
 Pipeline Mode
 
-Pipeline mode allows sending queries without having read the results of previously sent queries. It allows
-control of exactly how many and when network round trips occur.
+Pipeline mode allows sending queries without having read the results of previously sent queries. It allows control of
+exactly how many and when network round trips occur.
 
 Context Support
 
-All potentially blocking operations take a context.Context. If a context is canceled while the method is in progress the
-method immediately returns. In most circumstances, this will close the underlying connection.
+All potentially blocking operations take a context.Context. The default behavior when a context is canceled is for the
+method to immediately return. In most circumstances, this will also close the underlying connection. This behavior can
+be customized by using BuildContextWatcherHandler on the Config to create a ctxwatch.Handler with different behavior.
+This can be especially useful when queries that are frequently canceled and the overhead of creating new connections is
+a problem. DeadlineContextWatcherHandler and CancelRequestContextWatcherHandler can be used to introduce a delay before
+interrupting the query in such a way as to close the connection.
 
 The CancelRequest method may be used to request the PostgreSQL server cancel an in-progress query without forcing the
 client to abort.
diff --git a/pgconn/pgconn.go b/pgconn/pgconn.go
index 170f18db..f9930f34 100644
--- a/pgconn/pgconn.go
+++ b/pgconn/pgconn.go
@@ -18,8 +18,8 @@ import (
 
 	"github.com/jackc/pgx/v5/internal/iobufpool"
 	"github.com/jackc/pgx/v5/internal/pgio"
+	"github.com/jackc/pgx/v5/pgconn/ctxwatch"
 	"github.com/jackc/pgx/v5/pgconn/internal/bgreader"
-	"github.com/jackc/pgx/v5/pgconn/internal/ctxwatch"
 	"github.com/jackc/pgx/v5/pgproto3"
 )
 
@@ -82,6 +82,8 @@ type PgConn struct {
 	slowWriteTimer    *time.Timer
 	bgReaderStarted   chan struct{}
 
+	customData map[string]any
+
 	config *Config
 
 	status byte // One of connStatus* constants
@@ -278,31 +280,30 @@ func connect(ctx context.Context, config *Config, fallbackConfig *FallbackConfig
 	pgConn := new(PgConn)
 	pgConn.config = config
 	pgConn.cleanupDone = make(chan struct{})
+	pgConn.customData = make(map[string]any)
 
 	var err error
 	network, address := NetworkAddress(fallbackConfig.Host, fallbackConfig.Port)
-	netConn, err := config.DialFunc(ctx, network, address)
+	pgConn.conn, err = config.DialFunc(ctx, network, address)
 	if err != nil {
 		return nil, &ConnectError{Config: config, msg: "dial error", err: normalizeTimeoutError(ctx, err)}
 	}
 
-	pgConn.conn = netConn
-	pgConn.contextWatcher = newContextWatcher(netConn)
-	pgConn.contextWatcher.Watch(ctx)
-
 	if fallbackConfig.TLSConfig != nil {
-		nbTLSConn, err := startTLS(netConn, fallbackConfig.TLSConfig)
+		pgConn.contextWatcher = ctxwatch.NewContextWatcher(&DeadlineContextWatcherHandler{Conn: pgConn.conn})
+		pgConn.contextWatcher.Watch(ctx)
+		tlsConn, err := startTLS(pgConn.conn, fallbackConfig.TLSConfig)
 		pgConn.contextWatcher.Unwatch() // Always unwatch `netConn` after TLS.
 		if err != nil {
-			netConn.Close()
+			pgConn.conn.Close()
 			return nil, &ConnectError{Config: config, msg: "tls error", err: normalizeTimeoutError(ctx, err)}
 		}
 
-		pgConn.conn = nbTLSConn
-		pgConn.contextWatcher = newContextWatcher(nbTLSConn)
-		pgConn.contextWatcher.Watch(ctx)
+		pgConn.conn = tlsConn
 	}
 
+	pgConn.contextWatcher = ctxwatch.NewContextWatcher(config.BuildContextWatcherHandler(pgConn))
+	pgConn.contextWatcher.Watch(ctx)
 	defer pgConn.contextWatcher.Unwatch()
 
 	pgConn.parameterStatuses = make(map[string]string)
@@ -412,13 +413,6 @@ func connect(ctx context.Context, config *Config, fallbackConfig *FallbackConfig
 	}
 }
 
-func newContextWatcher(conn net.Conn) *ctxwatch.ContextWatcher {
-	return ctxwatch.NewContextWatcher(
-		func() { conn.SetDeadline(time.Date(1, 1, 1, 1, 1, 1, 1, time.UTC)) },
-		func() { conn.SetDeadline(time.Time{}) },
-	)
-}
-
 func startTLS(conn net.Conn, tlsConfig *tls.Config) (net.Conn, error) {
 	err := binary.Write(conn, binary.BigEndian, []int32{8, 80877103})
 	if err != nil {
@@ -988,10 +982,7 @@ func (pgConn *PgConn) CancelRequest(ctx context.Context) error {
 	defer cancelConn.Close()
 
 	if ctx != context.Background() {
-		contextWatcher := ctxwatch.NewContextWatcher(
-			func() { cancelConn.SetDeadline(time.Date(1, 1, 1, 1, 1, 1, 1, time.UTC)) },
-			func() { cancelConn.SetDeadline(time.Time{}) },
-		)
+		contextWatcher := ctxwatch.NewContextWatcher(&DeadlineContextWatcherHandler{Conn: cancelConn})
 		contextWatcher.Watch(ctx)
 		defer contextWatcher.Unwatch()
 	}
@@ -1880,6 +1871,11 @@ func (pgConn *PgConn) SyncConn(ctx context.Context) error {
 	return errors.New("SyncConn: conn never synchronized")
 }
 
+// CustomData returns a map that can be used to associate custom data with the connection.
+func (pgConn *PgConn) CustomData() map[string]any {
+	return pgConn.customData
+}
+
 // HijackedConn is the result of hijacking a connection.
 //
 // Due to the necessary exposure of internal implementation details, it is not covered by the semantic versioning
@@ -1892,6 +1888,7 @@ type HijackedConn struct {
 	TxStatus          byte
 	Frontend          *pgproto3.Frontend
 	Config            *Config
+	CustomData        map[string]any
 }
 
 // Hijack extracts the internal connection data. pgConn must be in an idle state. SyncConn should be called immediately
@@ -1914,6 +1911,7 @@ func (pgConn *PgConn) Hijack() (*HijackedConn, error) {
 		TxStatus:          pgConn.txStatus,
 		Frontend:          pgConn.frontend,
 		Config:            pgConn.config,
+		CustomData:        pgConn.customData,
 	}, nil
 }
 
@@ -1933,13 +1931,14 @@ func Construct(hc *HijackedConn) (*PgConn, error) {
 		txStatus:          hc.TxStatus,
 		frontend:          hc.Frontend,
 		config:            hc.Config,
+		customData:        hc.CustomData,
 
 		status: connStatusIdle,
 
 		cleanupDone: make(chan struct{}),
 	}
 
-	pgConn.contextWatcher = newContextWatcher(pgConn.conn)
+	pgConn.contextWatcher = ctxwatch.NewContextWatcher(hc.Config.BuildContextWatcherHandler(pgConn))
 	pgConn.bgReader = bgreader.New(pgConn.conn)
 	pgConn.slowWriteTimer = time.AfterFunc(time.Duration(math.MaxInt64),
 		func() {
@@ -2246,3 +2245,71 @@ func (p *Pipeline) Close() error {
 
 	return p.err
 }
+
+// DeadlineContextWatcherHandler handles canceled contexts by setting a deadline on a net.Conn.
+type DeadlineContextWatcherHandler struct {
+	Conn net.Conn
+
+	// DeadlineDelay is the delay to set on the deadline set on net.Conn when the context is canceled.
+	DeadlineDelay time.Duration
+}
+
+func (h *DeadlineContextWatcherHandler) HandleCancel(ctx context.Context) {
+	h.Conn.SetDeadline(time.Now().Add(h.DeadlineDelay))
+}
+
+func (h *DeadlineContextWatcherHandler) HandleUnwatchAfterCancel() {
+	h.Conn.SetDeadline(time.Time{})
+}
+
+// CancelRequestContextWatcherHandler handles canceled contexts by sending a cancel request to the server. It also sets
+// a deadline on a net.Conn as a fallback.
+type CancelRequestContextWatcherHandler struct {
+	Conn *PgConn
+
+	// CancelRequestDelay is the delay before sending the cancel request to the server.
+	CancelRequestDelay time.Duration
+
+	// DeadlineDelay is the delay to set on the deadline set on net.Conn when the context is canceled.
+	DeadlineDelay time.Duration
+
+	cancelFinishedChan             chan struct{}
+	handleUnwatchAfterCancelCalled func()
+}
+
+func (h *CancelRequestContextWatcherHandler) HandleCancel(context.Context) {
+	h.cancelFinishedChan = make(chan struct{})
+	var handleUnwatchedAfterCancelCalledCtx context.Context
+	handleUnwatchedAfterCancelCalledCtx, h.handleUnwatchAfterCancelCalled = context.WithCancel(context.Background())
+
+	deadline := time.Now().Add(h.DeadlineDelay)
+	h.Conn.conn.SetDeadline(deadline)
+
+	go func() {
+		defer close(h.cancelFinishedChan)
+
+		select {
+		case <-handleUnwatchedAfterCancelCalledCtx.Done():
+			return
+		case <-time.After(h.CancelRequestDelay):
+		}
+
+		cancelRequestCtx, cancel := context.WithDeadline(handleUnwatchedAfterCancelCalledCtx, deadline)
+		defer cancel()
+		h.Conn.CancelRequest(cancelRequestCtx)
+
+		// CancelRequest is inherently racy. Even though the cancel request has been received by the server at this point,
+		// it hasn't necessarily been delivered to the other connection. If we immediately return and the connection is
+		// immediately used then it is possible the CancelRequest will actually cancel our next query. The
+		// TestCancelRequestContextWatcherHandler Stress test can produce this error without the sleep below. The sleep time
+		// is arbitrary, but should be sufficient to prevent this error case.
+		time.Sleep(100 * time.Millisecond)
+	}()
+}
+
+func (h *CancelRequestContextWatcherHandler) HandleUnwatchAfterCancel() {
+	h.handleUnwatchAfterCancelCalled()
+	<-h.cancelFinishedChan
+
+	h.Conn.conn.SetDeadline(time.Time{})
+}
diff --git a/pgtype/pgtype_default.go b/pgtype/pgtype_default.go
index d56d95bd..38093ef4 100644
--- a/pgtype/pgtype_default.go
+++ b/pgtype/pgtype_default.go
@@ -82,8 +82,8 @@ func initDefaultMap() {
 	defaultMap.RegisterType(&Type{Name: "text", OID: TextOID, Codec: TextCodec{}})
 	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: "timestamp", OID: TimestampOID, Codec: &TimestampCodec{}})
+	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/timestamp.go b/pgtype/timestamp.go
index 35d73956..677a2c6e 100644
--- a/pgtype/timestamp.go
+++ b/pgtype/timestamp.go
@@ -46,7 +46,7 @@ func (ts *Timestamp) Scan(src any) error {
 
 	switch src := src.(type) {
 	case string:
-		return scanPlanTextTimestampToTimestampScanner{}.Scan([]byte(src), ts)
+		return (&scanPlanTextTimestampToTimestampScanner{}).Scan([]byte(src), ts)
 	case time.Time:
 		*ts = Timestamp{Time: src, Valid: true}
 		return nil
@@ -116,17 +116,21 @@ func (ts *Timestamp) UnmarshalJSON(b []byte) error {
 	return nil
 }
 
-type TimestampCodec struct{}
+type TimestampCodec struct {
+	// ScanLocation is the location that the time is assumed to be in for scanning. This is different from
+	// TimestamptzCodec.ScanLocation in that this setting does change the instant in time that the timestamp represents.
+	ScanLocation *time.Location
+}
 
-func (TimestampCodec) FormatSupported(format int16) bool {
+func (*TimestampCodec) FormatSupported(format int16) bool {
 	return format == TextFormatCode || format == BinaryFormatCode
 }
 
-func (TimestampCodec) PreferredFormat() int16 {
+func (*TimestampCodec) PreferredFormat() int16 {
 	return BinaryFormatCode
 }
 
-func (TimestampCodec) PlanEncode(m *Map, oid uint32, format int16, value any) EncodePlan {
+func (*TimestampCodec) PlanEncode(m *Map, oid uint32, format int16, value any) EncodePlan {
 	if _, ok := value.(TimestampValuer); !ok {
 		return nil
 	}
@@ -220,27 +224,27 @@ func discardTimeZone(t time.Time) time.Time {
 	return t
 }
 
-func (TimestampCodec) PlanScan(m *Map, oid uint32, format int16, target any) ScanPlan {
+func (c *TimestampCodec) PlanScan(m *Map, oid uint32, format int16, target any) ScanPlan {
 
 	switch format {
 	case BinaryFormatCode:
 		switch target.(type) {
 		case TimestampScanner:
-			return scanPlanBinaryTimestampToTimestampScanner{}
+			return &scanPlanBinaryTimestampToTimestampScanner{location: c.ScanLocation}
 		}
 	case TextFormatCode:
 		switch target.(type) {
 		case TimestampScanner:
-			return scanPlanTextTimestampToTimestampScanner{}
+			return &scanPlanTextTimestampToTimestampScanner{location: c.ScanLocation}
 		}
 	}
 
 	return nil
 }
 
-type scanPlanBinaryTimestampToTimestampScanner struct{}
+type scanPlanBinaryTimestampToTimestampScanner struct{ location *time.Location }
 
-func (scanPlanBinaryTimestampToTimestampScanner) Scan(src []byte, dst any) error {
+func (plan *scanPlanBinaryTimestampToTimestampScanner) Scan(src []byte, dst any) error {
 	scanner := (dst).(TimestampScanner)
 
 	if src == nil {
@@ -264,15 +268,18 @@ func (scanPlanBinaryTimestampToTimestampScanner) Scan(src []byte, dst any) error
 			microsecFromUnixEpochToY2K/1000000+microsecSinceY2K/1000000,
 			(microsecFromUnixEpochToY2K%1000000*1000)+(microsecSinceY2K%1000000*1000),
 		).UTC()
+		if plan.location != nil {
+			tim = time.Date(tim.Year(), tim.Month(), tim.Day(), tim.Hour(), tim.Minute(), tim.Second(), tim.Nanosecond(), plan.location)
+		}
 		ts = Timestamp{Time: tim, Valid: true}
 	}
 
 	return scanner.ScanTimestamp(ts)
 }
 
-type scanPlanTextTimestampToTimestampScanner struct{}
+type scanPlanTextTimestampToTimestampScanner struct{ location *time.Location }
 
-func (scanPlanTextTimestampToTimestampScanner) Scan(src []byte, dst any) error {
+func (plan *scanPlanTextTimestampToTimestampScanner) Scan(src []byte, dst any) error {
 	scanner := (dst).(TimestampScanner)
 
 	if src == nil {
@@ -302,13 +309,17 @@ func (scanPlanTextTimestampToTimestampScanner) Scan(src []byte, dst any) error {
 			tim = time.Date(year, tim.Month(), tim.Day(), tim.Hour(), tim.Minute(), tim.Second(), tim.Nanosecond(), tim.Location())
 		}
 
+		if plan.location != nil {
+			tim = time.Date(tim.Year(), tim.Month(), tim.Day(), tim.Hour(), tim.Minute(), tim.Second(), tim.Nanosecond(), plan.location)
+		}
+
 		ts = Timestamp{Time: tim, Valid: true}
 	}
 
 	return scanner.ScanTimestamp(ts)
 }
 
-func (c TimestampCodec) DecodeDatabaseSQLValue(m *Map, oid uint32, format int16, src []byte) (driver.Value, error) {
+func (c *TimestampCodec) DecodeDatabaseSQLValue(m *Map, oid uint32, format int16, src []byte) (driver.Value, error) {
 	if src == nil {
 		return nil, nil
 	}
@@ -326,7 +337,7 @@ func (c TimestampCodec) DecodeDatabaseSQLValue(m *Map, oid uint32, format int16,
 	return ts.Time, nil
 }
 
-func (c TimestampCodec) DecodeValue(m *Map, oid uint32, format int16, src []byte) (any, error) {
+func (c *TimestampCodec) DecodeValue(m *Map, oid uint32, format int16, src []byte) (any, error) {
 	if src == nil {
 		return nil, nil
 	}
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/rows.go b/rows.go
index 78723d6c..d4f7a901 100644
--- a/rows.go
+++ b/rows.go
@@ -419,6 +419,8 @@ type CollectableRow interface {
 type RowToFunc[T any] func(row CollectableRow) (T, error)
 
 // AppendRows iterates through rows, calling fn for each row, and appending the results into a slice of T.
+//
+// This function closes the rows automatically on return.
 func AppendRows[T any, S ~[]T](slice S, rows Rows, fn RowToFunc[T]) (S, error) {
 	defer rows.Close()
 
@@ -438,12 +440,16 @@ func AppendRows[T any, S ~[]T](slice S, rows Rows, fn RowToFunc[T]) (S, error) {
 }
 
 // CollectRows iterates through rows, calling fn for each row, and collecting the results into a slice of T.
+//
+// This function closes the rows automatically on return.
 func CollectRows[T any](rows Rows, fn RowToFunc[T]) ([]T, error) {
 	return AppendRows([]T{}, rows, fn)
 }
 
 // CollectOneRow calls fn for the first row in rows and returns the result. If no rows are found returns an error where errors.Is(ErrNoRows) is true.
 // CollectOneRow is to CollectRows as QueryRow is to Query.
+//
+// This function closes the rows automatically on return.
 func CollectOneRow[T any](rows Rows, fn RowToFunc[T]) (T, error) {
 	defer rows.Close()
 
@@ -469,6 +475,8 @@ func CollectOneRow[T any](rows Rows, fn RowToFunc[T]) (T, error) {
 // CollectExactlyOneRow calls fn for the first row in rows and returns the result.
 //   - If no rows are found returns an error where errors.Is(ErrNoRows) is true.
 //   - If more than 1 row is found returns an error where errors.Is(ErrTooManyRows) is true.
+//
+// This function closes the rows automatically on return.
 func CollectExactlyOneRow[T any](rows Rows, fn RowToFunc[T]) (T, error) {
 	defer rows.Close()
 


Test patch:
diff --git a/pgconn/internal/ctxwatch/context_watcher_test.go b/pgconn/ctxwatch/context_watcher_test.go
similarity index 66%
rename from pgconn/internal/ctxwatch/context_watcher_test.go
rename to pgconn/ctxwatch/context_watcher_test.go
index d62b97d3..302aabe3 100644
--- a/pgconn/internal/ctxwatch/context_watcher_test.go
+++ b/pgconn/ctxwatch/context_watcher_test.go
@@ -6,17 +6,32 @@ import (
 	"testing"
 	"time"
 
-	"github.com/jackc/pgx/v5/pgconn/internal/ctxwatch"
+	"github.com/jackc/pgx/v5/pgconn/ctxwatch"
 	"github.com/stretchr/testify/require"
 )
 
+type testHandler struct {
+	handleCancel             func(context.Context)
+	handleUnwatchAfterCancel func()
+}
+
+func (h *testHandler) HandleCancel(ctx context.Context) {
+	h.handleCancel(ctx)
+}
+
+func (h *testHandler) HandleUnwatchAfterCancel() {
+	h.handleUnwatchAfterCancel()
+}
+
 func TestContextWatcherContextCancelled(t *testing.T) {
 	canceledChan := make(chan struct{})
 	cleanupCalled := false
-	cw := ctxwatch.NewContextWatcher(func() {
-		canceledChan <- struct{}{}
-	}, func() {
-		cleanupCalled = true
+	cw := ctxwatch.NewContextWatcher(&testHandler{
+		handleCancel: func(context.Context) {
+			canceledChan <- struct{}{}
+		}, handleUnwatchAfterCancel: func() {
+			cleanupCalled = true
+		},
 	})
 
 	ctx, cancel := context.WithCancel(context.Background())
@@ -35,10 +50,12 @@ func TestContextWatcherContextCancelled(t *testing.T) {
 }
 
 func TestContextWatcherUnwatchdBeforeContextCancelled(t *testing.T) {
-	cw := ctxwatch.NewContextWatcher(func() {
-		t.Error("cancel func should not have been called")
-	}, func() {
-		t.Error("cleanup func should not have been called")
+	cw := ctxwatch.NewContextWatcher(&testHandler{
+		handleCancel: func(context.Context) {
+			t.Error("cancel func should not have been called")
+		}, handleUnwatchAfterCancel: func() {
+			t.Error("cleanup func should not have been called")
+		},
 	})
 
 	ctx, cancel := context.WithCancel(context.Background())
@@ -48,7 +65,7 @@ func TestContextWatcherUnwatchdBeforeContextCancelled(t *testing.T) {
 }
 
 func TestContextWatcherMultipleWatchPanics(t *testing.T) {
-	cw := ctxwatch.NewContextWatcher(func() {}, func() {})
+	cw := ctxwatch.NewContextWatcher(&testHandler{handleCancel: func(context.Context) {}, handleUnwatchAfterCancel: func() {}})
 
 	ctx, cancel := context.WithCancel(context.Background())
 	defer cancel()
@@ -61,7 +78,7 @@ func TestContextWatcherMultipleWatchPanics(t *testing.T) {
 }
 
 func TestContextWatcherUnwatchWhenNotWatchingIsSafe(t *testing.T) {
-	cw := ctxwatch.NewContextWatcher(func() {}, func() {})
+	cw := ctxwatch.NewContextWatcher(&testHandler{handleCancel: func(context.Context) {}, handleUnwatchAfterCancel: func() {}})
 	cw.Unwatch() // unwatch when not / never watching
 
 	ctx, cancel := context.WithCancel(context.Background())
@@ -72,7 +89,7 @@ func TestContextWatcherUnwatchWhenNotWatchingIsSafe(t *testing.T) {
 }
 
 func TestContextWatcherUnwatchIsConcurrencySafe(t *testing.T) {
-	cw := ctxwatch.NewContextWatcher(func() {}, func() {})
+	cw := ctxwatch.NewContextWatcher(&testHandler{handleCancel: func(context.Context) {}, handleUnwatchAfterCancel: func() {}})
 
 	ctx, cancel := context.WithTimeout(context.Background(), 100*time.Millisecond)
 	defer cancel()
@@ -88,10 +105,12 @@ func TestContextWatcherStress(t *testing.T) {
 	var cancelFuncCalls int64
 	var cleanupFuncCalls int64
 
-	cw := ctxwatch.NewContextWatcher(func() {
-		atomic.AddInt64(&cancelFuncCalls, 1)
-	}, func() {
-		atomic.AddInt64(&cleanupFuncCalls, 1)
+	cw := ctxwatch.NewContextWatcher(&testHandler{
+		handleCancel: func(context.Context) {
+			atomic.AddInt64(&cancelFuncCalls, 1)
+		}, handleUnwatchAfterCancel: func() {
+			atomic.AddInt64(&cleanupFuncCalls, 1)
+		},
 	})
 
 	cycleCount := 100000
@@ -134,7 +153,7 @@ func TestContextWatcherStress(t *testing.T) {
 }
 
 func BenchmarkContextWatcherUncancellable(b *testing.B) {
-	cw := ctxwatch.NewContextWatcher(func() {}, func() {})
+	cw := ctxwatch.NewContextWatcher(&testHandler{handleCancel: func(context.Context) {}, handleUnwatchAfterCancel: func() {}})
 
 	for i := 0; i < b.N; i++ {
 		cw.Watch(context.Background())
@@ -143,7 +162,7 @@ func BenchmarkContextWatcherUncancellable(b *testing.B) {
 }
 
 func BenchmarkContextWatcherCancelled(b *testing.B) {
-	cw := ctxwatch.NewContextWatcher(func() {}, func() {})
+	cw := ctxwatch.NewContextWatcher(&testHandler{handleCancel: func(context.Context) {}, handleUnwatchAfterCancel: func() {}})
 
 	for i := 0; i < b.N; i++ {
 		ctx, cancel := context.WithCancel(context.Background())
@@ -154,7 +173,7 @@ func BenchmarkContextWatcherCancelled(b *testing.B) {
 }
 
 func BenchmarkContextWatcherCancellable(b *testing.B) {
-	cw := ctxwatch.NewContextWatcher(func() {}, func() {})
+	cw := ctxwatch.NewContextWatcher(&testHandler{handleCancel: func(context.Context) {}, handleUnwatchAfterCancel: func() {}})
 
 	ctx, cancel := context.WithCancel(context.Background())
 	defer cancel()
diff --git a/pgconn/pgconn_test.go b/pgconn/pgconn_test.go
index 604aa9e6..40459214 100644
--- a/pgconn/pgconn_test.go
+++ b/pgconn/pgconn_test.go
@@ -24,6 +24,7 @@ import (
 	"github.com/jackc/pgx/v5/internal/pgio"
 	"github.com/jackc/pgx/v5/internal/pgmock"
 	"github.com/jackc/pgx/v5/pgconn"
+	"github.com/jackc/pgx/v5/pgconn/ctxwatch"
 	"github.com/jackc/pgx/v5/pgproto3"
 	"github.com/jackc/pgx/v5/pgtype"
 )
@@ -3184,6 +3185,22 @@ func TestConnOnPgError(t *testing.T) {
 	assert.True(t, pgConn.IsClosed())
 }
 
+func TestConnCustomData(t *testing.T) {
+	t.Parallel()
+
+	ctx, cancel := context.WithTimeout(context.Background(), 120*time.Second)
+	defer cancel()
+
+	pgConn, err := pgconn.Connect(ctx, os.Getenv("PGX_TEST_DATABASE"))
+	require.NoError(t, err)
+	defer closeConn(t, pgConn)
+
+	pgConn.CustomData()["foo"] = "bar"
+	assert.Equal(t, "bar", pgConn.CustomData()["foo"])
+
+	ensureConnValid(t, pgConn)
+}
+
 func Example() {
 	ctx, cancel := context.WithTimeout(context.Background(), 120*time.Second)
 	defer cancel()
@@ -3480,3 +3497,161 @@ func mustEncode(buf []byte, err error) []byte {
 	}
 	return buf
 }
+
+func TestDeadlineContextWatcherHandler(t *testing.T) {
+	t.Parallel()
+
+	t.Run("DeadlineExceeded with zero DeadlineDelay", func(t *testing.T) {
+		config, err := pgconn.ParseConfig(os.Getenv("PGX_TEST_DATABASE"))
+		require.NoError(t, err)
+		config.BuildContextWatcherHandler = func(conn *pgconn.PgConn) ctxwatch.Handler {
+			return &pgconn.DeadlineContextWatcherHandler{Conn: conn.Conn()}
+		}
+		config.ConnectTimeout = 5 * time.Second
+
+		pgConn, err := pgconn.ConnectConfig(context.Background(), config)
+		require.NoError(t, err)
+		defer closeConn(t, pgConn)
+
+		ctx, cancel := context.WithTimeout(context.Background(), 50*time.Millisecond)
+		defer cancel()
+
+		_, err = pgConn.Exec(ctx, "select 1, pg_sleep(1)").ReadAll()
+		require.Error(t, err)
+		require.ErrorIs(t, err, context.DeadlineExceeded)
+		require.True(t, pgConn.IsClosed())
+	})
+
+	t.Run("DeadlineExceeded with DeadlineDelay", func(t *testing.T) {
+		config, err := pgconn.ParseConfig(os.Getenv("PGX_TEST_DATABASE"))
+		require.NoError(t, err)
+		config.BuildContextWatcherHandler = func(conn *pgconn.PgConn) ctxwatch.Handler {
+			return &pgconn.DeadlineContextWatcherHandler{Conn: conn.Conn(), DeadlineDelay: 500 * time.Millisecond}
+		}
+		config.ConnectTimeout = 5 * time.Second
+
+		pgConn, err := pgconn.ConnectConfig(context.Background(), config)
+		require.NoError(t, err)
+		defer closeConn(t, pgConn)
+
+		ctx, cancel := context.WithTimeout(context.Background(), 100*time.Millisecond)
+		defer cancel()
+
+		_, err = pgConn.Exec(ctx, "select 1, pg_sleep(0.250)").ReadAll()
+		require.NoError(t, err)
+
+		ensureConnValid(t, pgConn)
+	})
+}
+
+func TestCancelRequestContextWatcherHandler(t *testing.T) {
+	t.Parallel()
+
+	t.Run("DeadlineExceeded cancels request after CancelRequestDelay", func(t *testing.T) {
+		config, err := pgconn.ParseConfig(os.Getenv("PGX_TEST_DATABASE"))
+		require.NoError(t, err)
+		config.BuildContextWatcherHandler = func(conn *pgconn.PgConn) ctxwatch.Handler {
+			return &pgconn.CancelRequestContextWatcherHandler{
+				Conn:               conn,
+				CancelRequestDelay: 250 * time.Millisecond,
+				DeadlineDelay:      5000 * time.Millisecond,
+			}
+		}
+		config.ConnectTimeout = 5 * time.Second
+
+		pgConn, err := pgconn.ConnectConfig(context.Background(), config)
+		require.NoError(t, err)
+		defer closeConn(t, pgConn)
+
+		ctx, cancel := context.WithTimeout(context.Background(), 50*time.Millisecond)
+		defer cancel()
+
+		_, err = pgConn.Exec(ctx, "select 1, pg_sleep(3)").ReadAll()
+		require.Error(t, err)
+		var pgErr *pgconn.PgError
+		require.ErrorAs(t, err, &pgErr)
+
+		ensureConnValid(t, pgConn)
+	})
+
+	t.Run("DeadlineExceeded - do not send cancel request when query finishes in grace period", func(t *testing.T) {
+		config, err := pgconn.ParseConfig(os.Getenv("PGX_TEST_DATABASE"))
+		require.NoError(t, err)
+		config.BuildContextWatcherHandler = func(conn *pgconn.PgConn) ctxwatch.Handler {
+			return &pgconn.CancelRequestContextWatcherHandler{
+				Conn:               conn,
+				CancelRequestDelay: 1000 * time.Millisecond,
+				DeadlineDelay:      5000 * time.Millisecond,
+			}
+		}
+		config.ConnectTimeout = 5 * time.Second
+
+		pgConn, err := pgconn.ConnectConfig(context.Background(), config)
+		require.NoError(t, err)
+		defer closeConn(t, pgConn)
+
+		ctx, cancel := context.WithTimeout(context.Background(), 100*time.Millisecond)
+		defer cancel()
+
+		_, err = pgConn.Exec(ctx, "select 1, pg_sleep(0.250)").ReadAll()
+		require.NoError(t, err)
+
+		ensureConnValid(t, pgConn)
+	})
+
+	t.Run("DeadlineExceeded sets conn deadline with DeadlineDelay", func(t *testing.T) {
+		config, err := pgconn.ParseConfig(os.Getenv("PGX_TEST_DATABASE"))
+		require.NoError(t, err)
+		config.BuildContextWatcherHandler = func(conn *pgconn.PgConn) ctxwatch.Handler {
+			return &pgconn.CancelRequestContextWatcherHandler{
+				Conn:               conn,
+				CancelRequestDelay: 5000 * time.Millisecond, // purposely setting this higher than DeadlineDelay to ensure the cancel request never happens.
+				DeadlineDelay:      250 * time.Millisecond,
+			}
+		}
+		config.ConnectTimeout = 5 * time.Second
+
+		pgConn, err := pgconn.ConnectConfig(context.Background(), config)
+		require.NoError(t, err)
+		defer closeConn(t, pgConn)
+
+		ctx, cancel := context.WithTimeout(context.Background(), 50*time.Millisecond)
+		defer cancel()
+
+		_, err = pgConn.Exec(ctx, "select 1, pg_sleep(1)").ReadAll()
+		require.Error(t, err)
+		require.ErrorIs(t, err, context.DeadlineExceeded)
+		require.True(t, pgConn.IsClosed())
+	})
+
+	for i := 0; i < 10; i++ {
+		t.Run(fmt.Sprintf("Stress %d", i), func(t *testing.T) {
+			t.Parallel()
+
+			config, err := pgconn.ParseConfig(os.Getenv("PGX_TEST_DATABASE"))
+			require.NoError(t, err)
+			config.BuildContextWatcherHandler = func(conn *pgconn.PgConn) ctxwatch.Handler {
+				return &pgconn.CancelRequestContextWatcherHandler{
+					Conn:               conn,
+					CancelRequestDelay: 5 * time.Millisecond,
+					DeadlineDelay:      1000 * time.Millisecond,
+				}
+			}
+			config.ConnectTimeout = 5 * time.Second
+
+			pgConn, err := pgconn.ConnectConfig(context.Background(), config)
+			require.NoError(t, err)
+			defer closeConn(t, pgConn)
+
+			for i := 0; i < 20; i++ {
+				func() {
+					ctx, cancel := context.WithTimeout(context.Background(), 4*time.Millisecond)
+					defer cancel()
+					pgConn.Exec(ctx, "select 1, pg_sleep(0.010)").ReadAll()
+					time.Sleep(100 * time.Millisecond) // ensure a cancel request that was a little late doesn't interrupt ensureConnValid.
+					ensureConnValid(t, pgConn)
+				}()
+			}
+		})
+	}
+}
diff --git a/pgtype/timestamp_test.go b/pgtype/timestamp_test.go
index 24f229d5..31b3ad82 100644
--- a/pgtype/timestamp_test.go
+++ b/pgtype/timestamp_test.go
@@ -38,6 +38,42 @@ func TestTimestampCodec(t *testing.T) {
 	})
 }
 
+func TestTimestampCodecWithScanLocationUTC(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:  "timestamp",
+			OID:   pgtype.TimestampOID,
+			Codec: &pgtype.TimestampCodec{ScanLocation: time.UTC},
+		})
+	}
+
+	pgxtest.RunValueRoundTripTests(context.Background(), t, connTestRunner, nil, "timestamp", []pgxtest.ValueRoundTripTest{
+		// Have to use pgtype.Timestamp instead of time.Time as source because otherwise the simple and exec query exec
+		// modes will encode the time for timestamptz. That is, they will convert it from local time zone.
+		{pgtype.Timestamp{Time: time.Date(2000, 1, 1, 0, 0, 0, 0, time.Local), Valid: true}, new(time.Time), isExpectedEq(time.Date(2000, 1, 1, 0, 0, 0, 0, time.UTC))},
+	})
+}
+
+func TestTimestampCodecWithScanLocationLocal(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:  "timestamp",
+			OID:   pgtype.TimestampOID,
+			Codec: &pgtype.TimestampCodec{ScanLocation: time.Local},
+		})
+	}
+
+	pgxtest.RunValueRoundTripTests(context.Background(), t, connTestRunner, nil, "timestamp", []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.Local))},
+	})
+}
+
 // https://github.com/jackc/pgx/v4/pgtype/pull/128
 func TestTimestampTranscodeBigTimeBinary(t *testing.T) {
 	defaultConnTestRunner.RunTest(context.Background(), t, func(ctx context.Context, t testing.TB, conn *pgx.Conn) {
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) {